BatchNorm2d bug/unintended behaviour

Valerio_Biscione: <p>Hi, I spotted an unintended behaviour or possibly a bug due to this fix<br>
<a href="https://github.com/pytorch/pytorch/issues" rel="noopener nofollow ugc">[BatchNorm] Unexpected behaviour with track_running_stats · Issue #37823 · pytorch/pytorch (github.com)</a></p>
<p>but before opening an issue I would like to make sure I am understanding this correctly.</p>
<p>In the past (before that commit), we could change the behaviour of <code>BatchNorm</code> after the module was created. That is, we could either run it in train mode (stats computed on the batch) or in eval mode (using <code>running_*</code>) depending on the value of <code>.track_running_stats</code> which could be changed even after module creation with:</p>
<pre><code class="lang-auto">for m in self.net.modules():
    if isinstance(m, torch.nn.BatchNorm2d):    
         m.track_running_stats = False
</code></pre>
<p>Indeed, I have used this trick in the past to always compute batch stats even in eval mode and it worked.</p>
<p>The newer commit doesn’t allow this behaviour. That’s because it will set the batch in “train” mode only if <code>.running_*</code> are set to None, which only happens when <code>track_running_stats</code> is set to <code>False</code> when first initializing the batch module, and not later on like in the snippets above.</p>
<p>The line that establishes whether it’s train or eval mode is this (in the earlier commit):<br>
<code> self.training or not self.track_running_stats,</code><br>
and this (in the newer commit):<br>
<code> (self.running_mean is None) and (self.running_var is None)</code></p>
<p>As you can see, in the newer commit, setting <code>self.track_running_stats</code> to <code>False</code> <em>after</em> initializing, the module won’t set the batch in train mode (<code>running_* is not None</code>), but it will set it in eval  mode using the default <code>running_mean</code> and <code>_variance</code>: 0, and 1, creating disastrous results.</p>
<p>Quick fix: set <code>track_running_stats=False</code> when the module is initialized, not later.</p>
<p>Is this worth fixing?</p>
<p>ps: I am using PyTorch 1.7.0</p>

ptrblck: <p>Could you explain, what exactly should be fixed a bit more?</p>
<aside class="quote no-group" data-username="Valerio_Biscione" data-post="1" data-topic="112822">
<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/valerio_biscione/48/22759_2.png" class="avatar"> Valerio_Biscione:</div>
<blockquote>
<p>In the past (before that commit), we could change the behaviour of <code>BatchNorm</code> after the module was created. That is, we could either run it in train mode (stats computed on the batch) or in eval mode (using <code>running_*</code>) depending on the value of <code>.track_running_stats</code> which could be changed even after module creation</p>
</blockquote>
</aside>
<p>This was indeed the case, but you could also use the same behavior by calling <code>.train()</code> and <code>.eval()</code> on the batchnorm layers.<br>
I guess your use case might be a bit more special in a way, which used the “hack” of changing the <code>track_running_stats</code> attribute?</p>