Struct GradientAccumulationCallback 

pub struct GradientAccumulationCallback { /* private fields */ }

Gradient Accumulation callback.

Simulates larger batch sizes by accumulating gradients over multiple mini-batches before updating parameters. This is useful when GPU memory is limited but you want to train with effectively larger batches.

Effective batch size = mini_batch_size * accumulation_steps

Implementations

impl GradientAccumulationCallback

pub fn new(accumulation_steps: usize) -> TrainResult<Self>

Create a new Gradient Accumulation callback.

Arguments

• accumulation_steps - Number of mini-batches to accumulate (e.g., 4, 8, 16)

pub fn accumulate( &mut self, gradients: &HashMap<String, Array<f64, Ix2>>, ) -> TrainResult<()>

Accumulate gradients.

pub fn should_update(&self) -> bool

Check if we should perform an optimizer step.

pub fn get_and_reset(&mut self) -> HashMap<String, Array<f64, Ix2>>

Get averaged accumulated gradients and reset.

Trait Implementations

impl Callback for GradientAccumulationCallback

fn on_epoch_begin( &mut self, _epoch: usize, _state: &TrainingState, ) -> TrainResult<()>

Called at the beginning of an epoch.

fn on_train_begin(&mut self, _state: &TrainingState) -> TrainResult<()>

Called at the beginning of training.

fn on_train_end(&mut self, _state: &TrainingState) -> TrainResult<()>

Called at the end of training.

fn on_epoch_end( &mut self, _epoch: usize, _state: &TrainingState, ) -> TrainResult<()>

Called at the end of an epoch.

fn on_batch_begin( &mut self, _batch: usize, _state: &TrainingState, ) -> TrainResult<()>

Called at the beginning of a batch.

fn on_batch_end( &mut self, _batch: usize, _state: &TrainingState, ) -> TrainResult<()>

Called at the end of a batch.

fn on_validation_end(&mut self, _state: &TrainingState) -> TrainResult<()>

Called after validation.

fn should_stop(&self) -> bool

Check if training should stop early.