Struct DataParallelTrainer 

pub struct DataParallelTrainer<M: Model<Input = Tensor, Output = Tensor>> { /* private fields */ }

Data parallel trainer that wraps a model for distributed training

Implementations

impl<M: Model<Input = Tensor, Output = Tensor>> DataParallelTrainer<M>

pub fn new( model: M, process_group: Arc<dyn ProcessGroup>, config: DistributedConfig, ) -> Result<Self>

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

Forward pass through the model

pub fn backward(&self, gradients: &mut HashMap<String, Tensor>) -> Result<()>

Backward pass with gradient synchronization

pub fn broadcast_parameters(&self) -> Result<()>

Broadcast model parameters from rank 0 to all other ranks

pub fn model(&self) -> Arc<Mutex<M>>

Get the wrapped model

pub fn process_group(&self) -> Arc<dyn ProcessGroup>

Get process group