Struct DataParallel 

pub struct DataParallel { /* private fields */ }

Data Parallel model wrapper for single-node multi-GPU training

Implementations

impl DataParallel

pub fn new( model: Box<dyn Model<f32>>, devices: Vec<Device>, comm_runtime: Arc<RwLock<CommunicationRuntime>>, ) -> Result<Self>

Create new DataParallel model wrapper

pub fn forward_parallel( &self, inputs: &[Tensor<f32>], ) -> Result<Vec<Tensor<f32>>>

Perform forward pass with data parallelism

pub fn sync_gradients(&mut self) -> Result<()>

Synchronize gradients across all device replicas

pub fn set_sync_mode(&mut self, mode: SynchronizationMode)

Set synchronization mode

pub fn devices(&self) -> &[Device]

Get device replicas

Trait Implementations

impl Model<f32> for DataParallel

fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>>

fn parameters(&self) -> Vec<&Tensor<f32>>

fn parameters_mut(&mut self) -> Vec<&mut Tensor<f32>>

fn set_training(&mut self, training: bool)

fn zero_grad(&mut self)

fn as_any(&self) -> &dyn Any

Provide access to Any for downcasting

fn extract_features(&self, input: &Tensor<T>) -> Result<Option<Vec<Tensor<T>>>>

Extract intermediate features for knowledge distillation Returns None if the model doesn’t support feature extraction