Struct EMA 

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

Exponential Moving Average of model parameters.

Maintains a shadow copy that is a smoothed version of training parameters.

Example

let mut ema = EMA::new(model.parameters(), 0.999);

// Training loop:
optimizer.step(&grads)?;
ema.update(&model.parameters())?;

// Evaluation:
ema.apply()?;             // Write EMA weights into model
let output = model.forward(input)?;
ema.restore()?;           // Restore training weights

Implementations

impl<B> EMA<B>

where B: Backend,

pub fn new(params: Vec<Tensor<B>>, decay: f64) -> Result<EMA<B>, Error>

Create a new EMA tracker.

Arguments

• params: The model parameters to track
• decay: Decay rate (typical: 0.999 or 0.9999)

pub fn update(&mut self, current_params: &[Tensor<B>]) -> Result<(), Error>

Update the EMA shadow parameters with current model parameters.

Call this after each optimizer step.

pub fn update_with_warmup( &mut self, current_params: &[Tensor<B>], ) -> Result<(), Error>

Update using an adjusted decay that ramps up during early training.

The effective decay is: min(decay, (1 + num_updates) / (10 + num_updates)) This prevents the EMA from being too biased toward initial values.

pub fn apply(&mut self) -> Result<(), Error>

Apply EMA parameters to the model (for evaluation).

This saves the current training parameters so they can be restored with restore().

pub fn restore(&mut self) -> Result<(), Error>

Restore training parameters after apply().

pub fn decay(&self) -> f64

Get the decay rate.

pub fn num_updates(&self) -> u64

Get the number of updates performed.

pub fn shadow_values(&self, index: usize) -> &[f64]

Get the shadow (EMA) values for a specific parameter index.