Struct SWACallback 

pub struct SWACallback { /* private fields */ }

SWA (Stochastic Weight Averaging) callback.

Averages model parameters over the course of training, typically starting from a later epoch. This often leads to better generalization and wider optima.

Reference: Izmailov et al. “Averaging Weights Leads to Wider Optima and Better Generalization” (UAI 2018)

Implementations

impl SWACallback

pub fn new(start_epoch: usize, update_frequency: usize, verbose: bool) -> Self

Create a new SWA callback.

Arguments

• start_epoch - Epoch to start averaging (e.g., 0.75 * total_epochs)
• update_frequency - Average parameters every N epochs (typically 1)
• verbose - Whether to print progress

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

Update SWA parameters with current model parameters.

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

Get the SWA parameters.

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

Apply SWA parameters to the model.

pub fn is_ready(&self) -> bool

Check if SWA has collected any averages.

Trait Implementations

impl Callback for SWACallback

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

Called at the end of an epoch.

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

Called at the end of training.

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

Called at the beginning of training.

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

Called at the beginning 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.