Struct GradientAccumulationCallback 

pub struct GradientAccumulationCallback { /* private fields */ }

Gradient Accumulation callback with advanced features.

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

Features

• Memory-efficient in-place accumulation
• Multiple scaling strategies
• Gradient overflow detection
• Memory usage tracking
• Automatic gradient zeroing

Example

use tensorlogic_train::{GradientAccumulationCallback, GradientScalingStrategy};

let mut grad_accum = GradientAccumulationCallback::new(
    4, // accumulate over 4 mini-batches
    GradientScalingStrategy::Average,
).unwrap();

Implementations

impl GradientAccumulationCallback

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

Create a new Gradient Accumulation callback with default average scaling.

Arguments

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

pub fn with_strategy( accumulation_steps: usize, scaling_strategy: GradientScalingStrategy, ) -> TrainResult<Self>

Create a new Gradient Accumulation callback with specified scaling strategy.

Arguments

• accumulation_steps - Number of mini-batches to accumulate
• scaling_strategy - How to scale accumulated gradients

pub fn with_grad_clipping(self, max_norm: f64) -> Self

Enable gradient clipping during accumulation.

Arguments

• max_norm - Maximum gradient norm before clipping

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

Accumulate gradients with optional clipping and overflow detection.

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 scaled accumulated gradients and reset state.

pub fn get_stats(&self) -> GradientAccumulationStats

Get statistics about gradient accumulation.

pub fn reset(&mut self)

Reset all state without returning gradients (useful for error recovery).

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.