How to apply exponential moving average decay for variables?

username: jef
cooked: <p>I am reading following paper. And it uses EMA decay for variables.<br>
<a href="https://arxiv.org/pdf/1611.01603.pdf" rel="noopener nofollow ugc">BI-DIRECTIONAL ATTENTION FLOW FOR MACHINE COMPREHENSION</a></p>
<blockquote>
<p>During training, the moving averages of all weights of the model are<br>
maintained with the exponential decay rate of 0.999.</p>
</blockquote>
<p>They use TensorFlow and I found the related code of EMA.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/allenai/bi-att-flow/blob/master/basic/model.py#L229">
  <header class="source">

      <a href="https://github.com/allenai/bi-att-flow/blob/master/basic/model.py#L229" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/allenai/bi-att-flow/blob/master/basic/model.py#L229" target="_blank" rel="noopener nofollow ugc">allenai/bi-att-flow/blob/master/basic/model.py#L229</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="219" style="counter-reset: li-counter 218 ;">
          <li>    for var in tf.get_collection("ema/scalar", scope=self.scope):</li>
          <li>        ema_var = ema.average(var)</li>
          <li>        tf.scalar_summary(ema_var.op.name, ema_var)</li>
          <li>    for var in tf.get_collection("ema/vector", scope=self.scope):</li>
          <li>        ema_var = ema.average(var)</li>
          <li>        tf.histogram_summary(ema_var.op.name, ema_var)</li>
          <li></li>
          <li>    with tf.control_dependencies([ema_op]):</li>
          <li>        self.loss = tf.identity(self.loss)</li>
          <li></li>
          <li class="selected">def _build_var_ema(self):</li>
          <li>    self.var_ema = tf.train.ExponentialMovingAverage(self.config.var_decay)</li>
          <li>    ema = self.var_ema</li>
          <li>    ema_op = ema.apply(tf.trainable_variables())</li>
          <li>    with tf.control_dependencies([ema_op]):</li>
          <li>        self.loss = tf.identity(self.loss)</li>
          <li></li>
          <li>def get_loss(self):</li>
          <li>    return self.loss</li>
          <li></li>
          <li>def get_global_step(self):</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>In PyTorch, how do I apply EMA to Variables? In TensorFlow, there is <code>tf.train.ExponentialMovingAverage</code> class.<br>
<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_averages</a></p>

username: alexis-jacq
cooked: <pre><code class="lang-python">class EMA(nn.Module):
    def __init__(self, mu):
        super(EMA, self).__init__()
        self.mu = mu
        
    def forward(self,x, last_average):
        new_average = self.mu*x + (1-self.mu)*last_average
        return new_average

ema = EMA(0.999)
x = Variable(torch.rand(5),requires_grad=True)
average = Variable(torch.zeros(5),requires_grad=True)
average = ema(x, average)
</code></pre>

username: jef
cooked: <p><a class="mention" href="/u/alexis-jacq">@alexis-jacq</a> Thank you!<br>
I updated your code like below. I added shadow copy mechanism to save last value as average data.<br>
Does it make sense in PyTorch manner?</p>
<pre><code class="lang-auto">   class EMA(nn.Module):
       def __init__(self, mu):
           super(EMA, self).__init__()
           self.mu = mu
           self.shadow = {}

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

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

   ema = EMA(0.999)
   for name, param in model.named_parameters():
       if param.requires_grad:
           ema.register(name, param.data)

  # in batch training loop
  # for batch in batches:
       optimizer.step()
       for name, param in model.named_parameters():
           if param.requires_grad:
                param.data = ema(name, param.data)
</code></pre>

username: alexis-jacq
cooked: <p>Ok I see. In that case you don’t need a module inheritage, since you don’t want to take into account the EMA correction of your parameters while computing the next gradient. So you could simply have something like:</p>
<pre><code class="lang-python">class EMA():
       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 = self.mu * x + (1.0 - self.mu) * self.shadow[name]
           self.shadow[name] = new_average.clone()
           return new_average

   ema = EMA(0.999)
   for name, param in model.named_parameters():
       if param.requires_grad:
           ema.register(name, param.data)

  # in batch training loop
  # for batch in batches:
       optimizer.step()
       for name, param in model.named_parameters():
           if param.requires_grad:
                param.data = ema(name, param.data)
</code></pre>

username: ronghanghu
cooked: <p>As far as I understand, it should be</p>
<pre><code class="lang-auto">new_average = (1.0 - self.mu) * x + self.mu * self.shadow[name]
</code></pre>
<p>instead, since mu is 0.999. You want to multiply 0.999 on the running average, not the input x</p>

username: alexis-jacq
cooked: <p>My mistake, you are right. Unfortunatly – I don’t know why – I can’t edit my previous post.</p>

username: yucoian
cooked: <p><a class="mention" href="/u/smth">@smth</a><br>
<a class="mention" href="/u/apaszke">@apaszke</a><br>
please help us! Thanks a lot!</p>

username: Knurpsbram
cooked: <p>I am confused as to how to apply EMA. Is the shadow model or the actual model used to process training items and compute gradients? If its the shadow model, then I don’t see how it’s different from using a smaller learning rate with the optimizer. Is it the case that the shadow model is only used at test time?</p>

username: eason_long
cooked: <p>that’s right！ there are so many wrong implements!</p>

username: amsword
cooked: <p>in the implementation, the moving averaged results will be used for the next iterations (last sentence). Another potential solution is only to track the moving average, but the parameters in the network is still the results from optimizer. That is, only run ema(name, param.data), but does not assign it back to param.data.</p>
<p>Just wondering which one is better, or which one is the commonly used one?</p>

username: schow
cooked: <p>Is there a reason why we are using:</p>
<pre><code class="lang-auto">self.shadow[name] = new_average.clone()
</code></pre>
<p>instead of</p>
<pre><code class="lang-auto">self.shadow[name] = new_average.detach().clone()
</code></pre>
<p>since we really don’t care about propagating the gradients to the shadow copies of the parameters.</p>