Struct DataParallelTrainer 

pub struct DataParallelTrainer { /* private fields */ }

Data parallel trainer

Implementations

impl DataParallelTrainer

pub fn new(config: DistributedConfig) -> Self

pub fn split_data(&self, data: &Tensor, labels: &Tensor) -> (Tensor, Tensor)

Split data across workers

pub fn accumulate_gradients(&mut self, gradients: Vec<Vec<f32>>)

Accumulate local gradients

pub fn sync_gradients(&mut self) -> Vec<Vec<f32>>

Synchronize gradients across workers

pub fn all_reduce(&self, gradients: &[Vec<f32>]) -> Vec<Vec<f32>>

All-reduce operation for gradients

pub fn broadcast_parameters(&self, parameters: &[Vec<f32>]) -> Vec<Vec<f32>>

Broadcast parameters from rank 0 to all workers

pub fn is_master(&self) -> bool