Semi-supervised-learning
Semi-supervised-learning copied to clipboard
Published
20 hours ago •
microsoft
microsoft
Improve EMA module
🚀 Feature
Improvement of current EMA module should be made, for saving GPU memory and easier application of ema model.
Motivation
Current EMA model are on the same device of model: https://github.com/microsoft/Semi-supervised-learning/blob/74bed89b78580e1b083b922f48c3b08bc63ec9db/train.py#L253 Besides EMA update is separated from ema_model: https://github.com/microsoft/Semi-supervised-learning/blob/74bed89b78580e1b083b922f48c3b08bc63ec9db/semilearn/core/utils/misc.py#L131
Should make this ema update more flexible and save GPU memory
Alternatives
Timm implementation as a good reference: https://github.com/rwightman/pytorch-image-models/blob/a520da9b495422bc773fb5dfe10819acb8bd7c5c/timm/utils/model_ema.py#L82
