Exponential Moving Average in PyTorch

username: Toby
cooked: <p>Hi,<br>
Exponential Moving Average (EMA) is an important feature in state-of-the-art research, in Tensorflow they already implemented it with <a href="https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage?version=stable" rel="noopener nofollow ugc">tf.train.ExponentialMovingAverage</a>. I wonder why the Pytorch team has not released an official version of EMA.</p>
<p>In others related question, there is no expert confirm that it is the correct implementation:</p>
<aside class="quote quote-modified" data-post="1" data-topic="20631">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bang_liu/48/3339_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/exponential-moving-average-in-pytorch-for-weights-and-gradients/20631">Exponential Moving Average in PyTorch, for weights and gradients</a> <a class="badge-category__wrapper " href="/c/nlp/8"><span data-category-id="8" data-drop-close="true" class="badge-category " title="Topics related to Natural Language Processing"><span class="badge-category__name">nlp</span></span></a>
  </div>
  <blockquote>
    Do we need to apply exponential moving average to weights during training when we use Adam (or other optimizers)? 
My EMA is defined as: 
class EMA(object):

def __init__(self, mu):
    self.mu = mu
    self.shadow = {}

def register(self, name, val):
    self.shadow[name] = val.clone()

def __call__(self, name, x):
    assert name in self.shadow
    new_average = (1.0 - self.mu) * x + self.mu * self.shadow[name]
    self.shadow[name] = new_average.clone()
    return new_average

My code is like…
  </blockquote>
</aside>

<aside class="quote quote-modified" data-post="1" data-topic="10856">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jef/48/2233_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-apply-exponential-moving-average-decay-for-variables/10856">How to apply exponential moving average decay for variables?</a> 
  </div>
  <blockquote>
    I am reading following paper. And it uses EMA decay for variables. 
<a href="https://arxiv.org/pdf/1611.01603.pdf" rel="noopener nofollow ugc">BI-DIRECTIONAL ATTENTION FLOW FOR MACHINE COMPREHENSION</a> 

During training, the moving averages of all weights of the model are 
maintained with the exponential decay rate of 0.999. 

They use TensorFlow and I found the related code of EMA. 
 
 
In PyTorch, how do I apply EMA to Variables? In TensorFlow, there is tf.train.ExponentialMovingAverage class. 
<a href="https://www.tensorflow.org/versions/r0.12/api_docs/python/train/moving_averages" class="onebox" target="_blank" rel="noopener nofollow ugc">https://www.tensorflow.org/versions/r0.12/api_docs/python/train/moving_avera…</a>
  </blockquote>
</aside>

<p>Or in <a href="https://github.com/fadel/pytorch_ema" rel="noopener nofollow ugc">this repo</a>, the issue said there was some problem with the code.</p>