How could I fis bn moving average updates in some conditions but not fix it in others

username: coincheung
cooked: <p>Hi,</p>
<p>Suppose there is a case like this:</p>
<pre><code class="lang-auto">x1, x2, lb1, lb2 = dataloader.next
logits1 = model(x1)
logits2 = model(x2)
loss = criteria1(logits1, lb1) + criteria2(logits2, lb2)
...
</code></pre>
<p>The problem is that I need to update bn moving average status only on <code>x1</code>, and I do not need to update the moving average status on other forward computations. How could I do this with pytorch?</p>

username: ptrblck
cooked: <p>You could call <code>.eval()</code> on all batch norm layers in your model after passing <code>x1</code> to the model and before using <code>x2</code>. After it just reset these layers to <code>.train()</code> again.</p>
<p>If you don’t use any dropout layers, you could also just call <code>model.eval()/.train()</code>.</p>

username: coincheung
cooked: <p>thanks for replying!!</p>
<p>If I call <code>model.eval()</code>, the input tensor would be normalized with the running mean and running var rather than the batch statistics, can this be avoided ?</p>

username: ptrblck
cooked: <p>You could disable the running stats completely by setting <code>track_running_stats=False</code>.<br>
However, your use case seems to be:</p>
<ul>
<li>for <code>x1</code> use batch stats and update running estimates</li>
<li>for <code>x2</code> just use batch stats</li>
</ul>
<p>What about the affine parameters?<br>
Should they also be updated using the loss from <code>x2</code> or just <code>x1</code>?<br>
In the latter case, you could use two different batchnorm layers, pass a flag to <code>forward</code>, and switch between these layers depending if <code>x1</code> or <code>x2</code> was passed.</p>

username: coincheung
cooked: <p>Thanks for replying!!</p>
<p>The affine parameters are trained from both x1 and x2. Two different batchnorm layers can solve the running estimates problem, but the affine parameters are not shared between these two batchnorm layers in this way. Any suggestions ?</p>

username: ptrblck
cooked: <p>You could use a hacky way of setting <code>track_running_stats=False</code> for the <code>x2</code> input and reset the running stats manually.<br>
Here is a small example:</p>
<pre><code class="lang-python">bn = nn.BatchNorm2d(3, track_running_stats=True)
bn.eval()
print(bn.running_mean) # zeros
print(bn.running_var) # ones

x = torch.randn(2, 3, 4, 4)
out = bn(x)
print(out.mean()) # should NOT be perfectly normal, since running stats used
print(out.std())
print(bn.running_mean)  # not updated
print(bn.running_var)

# Disable running_states
# internally buffers are still valid, so we need to reset them manually
bn.track_running_stats = False
out2 = bn(x)
print(out2.mean()) # should be normal now
print(out2.std())
print(bn.running_mean) # unfortunately updated
print(bn.running_var)

# Reset running stats
with torch.no_grad():
    bn.running_mean = (bn.running_mean - x.mean([0, 2, 3]) * bn.momentum) / (1 - bn.momentum)
    bn.running_var =  (bn.running_var - x.var([0, 2, 3]) * bn.momentum) / (1 - bn.momentum)
print(bn.running_mean) # back to values before last update
print(bn.running_var)

# enable running stats again
bn.track_running_stats = True
out3 = bn(x)
print(out3.mean() == out.mean()) # compare to initial output
print(out3.std() == out.std())
</code></pre>
<p>I would recommend to test this approach in your model and make sure you’ll get the desired outputs, gradients and updates.</p>

username: coincheung
cooked: <p>Thanks a lot, I would try.</p>