EMA

class torch_uncertainty.models.EMA(core_model, momentum)

Exponential Moving Average (EMA).

The model given as argument is used to compute the gradient during the training. The EMA model is regularly updated with the inner-model and used at evaluation time.

Parameters:

• core_model (nn.Module) – The model to train and ensemble.

• momentum (float) – The momentum of the moving average. The larger the momentum, the more stable the model.

Note

The momentum value is often large, such as 0.9 or 0.95.

update_wrapper(epoch=None)

Update the EMA model.

Parameters:

epoch (int) – The current epoch. For API consistency.