Struct ModelEMACallback 

pub struct ModelEMACallback { /* private fields */ }

Model EMA (Exponential Moving Average) callback.

Maintains an exponential moving average of model parameters during training. This often leads to better generalization and more stable predictions.

The shadow parameters are updated as: shadow_param = decay * shadow_param + (1 - decay) * param

Reference: Common practice in modern deep learning, popularized by Mean Teacher and other semi-supervised learning methods.

Implementations

impl ModelEMACallback

pub fn new(decay: f64, use_warmup: bool) -> Self

Create a new Model EMA callback.

Arguments

• decay - EMA decay rate (e.g., 0.999, 0.9999)
• use_warmup - Whether to use decay warmup (recommended)

pub fn initialize(&mut self, parameters: &HashMap<String, Array<f64, Ix2>>)

Initialize shadow parameters from current model parameters.

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

Update EMA parameters.

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

Get the EMA parameters.

pub fn apply_shadow(&self, parameters: &mut HashMap<String, Array<f64, Ix2>>)

Apply EMA parameters to the model (for evaluation).

Trait Implementations

impl Callback for ModelEMACallback

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

Called at the beginning of training.

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

Called at the end of a batch.

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

Called at the end of training.

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

Called at the beginning of an epoch.

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_validation_end(&mut self, _state: &TrainingState) -> TrainResult<()>

Called after validation.

fn should_stop(&self) -> bool

Check if training should stop early.