Exponential Moving Average (EMA)

[miraodasilva]

🚀 Feature

Keep an Exponential Moving Average (EMA) of the model's weights as it is training. This is available on tensorflow but not on pytorch. https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage .

Motivation

EMA has shown to be extremely benefitial for bootstrapping better models from scratch. For instance efficientnet (v1 and v2) benefits heavily from the usage of this method. It's also widely used in Self-supervised Learning.

Pitch

Basically what needs to be done is already layed out here https://forums.pytorchlightning.ai/t/adopting-exponential-moving-average-ema-for-pl-pipeline/488 . This code requires this package https://github.com/fadel/pytorch_ema .

[justusschock]

I think having it as a callback would be nice.

@miraodasilva Are you willing to contribute this?

[miraodasilva]

Sorry, I don't really have the time do it properly right now. However, I will start working with it on my own a bit, and perhaps I will contribute in the future. Thanks for responding!

[tchaton]

Hey @miraodasilva,

It can be done using the Stochastic Weight Averaging Callback and replacing the avg_gn function there:
https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/stochastic_weight_avg.py#L44

Here is the mean average implementation:
https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/stochastic_weight_avg.py#L287

[miraodasilva]

I see, hadn't seen that, thanks a lot!

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[hal-314]

Only for future readers, I don't think that the given solution is equivalent. avg_fn is called once per epoch while EMA updates happens every training step. I don't think that EMA can be implemented with SWA callback, See #10914 for code.