Struct ParallelTrainer 

pub struct ParallelTrainer<M, O, B: Backend> {
    pub model: M,
    pub optimizer: O,
    /* private fields */
}

High-level training loop with gradient accumulation.

Splits a large effective batch into accumulation_steps micro-batches, accumulates gradients across all of them, then performs a single optimizer step. This simulates a larger batch size without requiring more memory.

Example

let model = Sequential::new(vec![...]);
let optimizer = Adam::new(model.parameters(), 1e-3);
let mut trainer = ParallelTrainer::new(model, optimizer, 4);

// Each call accumulates 1/4 of the gradient; every 4th call steps.
for (i, (x, y)) in data.iter().enumerate() {
    if let Some(loss) = trainer.accumulate_step(&x, &y, mse_loss)? {
        println!("step {}: loss = {:.4}", i, loss);
    }
}

Fields

model: M

The model.

optimizer: O

The optimizer.

Implementations

impl<M, O, B> ParallelTrainer<M, O, B>

where M: Module<B>, O: Optimizer<B>, B: Backend,

pub fn new(model: M, optimizer: O, accumulation_steps: usize) -> Self

Create a new ParallelTrainer.

accumulation_steps: number of micro-batches before each optimizer step.

pub fn accumulate_step<F>( &mut self, input: &Tensor<B>, target: &Tensor<B>, loss_fn: F, ) -> Result<Option<f64>>

where F: Fn(&Tensor<B>, &Tensor<B>) -> Result<Tensor<B>>,

Process one micro-batch. Returns Some(avg_loss) when an optimizer step was performed (every accumulation_steps calls), else None.

pub fn flush(&mut self) -> Result<Option<f64>>

Force an optimizer step with whatever gradients have been accumulated so far. Useful at the end of an epoch when remaining micro-batches < accumulation_steps.