Struct GradAccumulator 

pub struct GradAccumulator<B>
where
    B: Backend,
{ /* private fields */ }

Gradient accumulation helper.

Accumulates gradients over multiple mini-batches and provides the averaged gradient for the optimizer step.

Example

let mut accum = GradAccumulator::<CpuBackend>::new(4); // 4 accumulation steps

for (i, batch) in batches.iter().enumerate() {
    let loss = model.forward(batch)?;
    let grads = loss.backward()?;
     
    if let Some(avg_grads) = accum.step(&grads, &params)? {
        optimizer.step(&avg_grads)?;
    }
}

Implementations

impl<B> GradAccumulator<B>

where B: Backend,

pub fn new(accum_steps: u64) -> GradAccumulator<B>

Create a new gradient accumulator.

Arguments

• accum_steps: Number of mini-batches to accumulate before stepping

pub fn step( &mut self, grads: &GradStore<B>, params: &[Tensor<B>], ) -> Result<Option<GradStore<B>>, Error>

Add gradients from one mini-batch.

Returns Some(averaged_grads) when accum_steps batches have been accumulated, otherwise returns None.

The returned gradients are divided by accum_steps to produce the average gradient.

pub fn reset(&mut self)

Reset the accumulator, discarding any accumulated gradients.

pub fn accum_steps(&self) -> u64

Get the number of accumulation steps.

pub fn current_step(&self) -> u64

Get the current step within the accumulation window.