MomentumAnnealingEMA — mmengine 0.10.7 documentation (original) (raw)

class mmengine.model.MomentumAnnealingEMA(model, momentum=0.0002, gamma=100, interval=1, device=None, update_buffers=False)[source]¶

Exponential moving average (EMA) with momentum annealing strategy.

Parameters:

• model (nn.Module) – The model to be averaged.
• momentum (float) – The momentum used for updating ema parameter. Defaults to 0.0002. Ema’s parameter are updated with the formula\(averaged\_param = (1-momentum) * averaged\_param + momentum * source\_param\).
• gamma (int) – Use a larger momentum early in training and gradually annealing to a smaller value to update the ema model smoothly. The momentum is calculated as max(momentum, gamma / (gamma + steps)) Defaults to 100.
• interval (int) – Interval between two updates. Defaults to 1.
• device (torch.device, optional) – If provided, the averaged model will be stored on the device. Defaults to None.
• update_buffers (bool) – if True, it will compute running averages for both the parameters and the buffers of the model. Defaults to False.

avg_func(averaged_param, source_param, steps)[source]¶

Compute the moving average of the parameters using the linear momentum strategy.

Parameters:

• averaged_param (Tensor) – The averaged parameters.
• source_param (Tensor) – The source parameters.
• steps (int) – The number of times the parameters have been updated.

Return type:

None