Affine and momentum of BatchNorm layer

Mirsadeghi: <p>Hi,<br>
In a specific application, I need to freeze running statistic calculation in BatchNorm layer in a part of my code, but I need to utilize “gamma(weight)” and “beta(bias)” of this layer in training with gradient forward\backward pass.<br>
I have implemented this by building an extra BatchNorm layer with “affine” false and doing forward pass as:</p>
<p>base_BatchNorm = nn.BatchNorm1d(1200)<br>
extra_BatchNorm = nn.BatchNorm1d(1200, affine=False)<br>
x = extra_BatchNorm(x)*base_BatchNorm.weight+base_BatchNorm.bias</p>
<p>I am curiose if there is more clean way to implement this. I think this can be done by setting momentum=0? or by a combination of  momentum=0 and affine=False for base_BatchNorm  as:</p>
<p>base_BatchNorm.momentum = 0<br>
base_BatchNorm.affine = false<br>
x = base_BatchNorm(x)<br>
base_BatchNorm.momentum = 0.1<br>
base_BatchNorm.affine = true</p>
<p>Is it right?</p>

ptrblck: <p>You can call <code>bn.eval()</code> which will use the running stats and will not update them. The affine parameters will still be trained.</p>

Mirsadeghi: <p>I don’t want to use running_stats in this part of my code. I want to apply zero-mean, unit_variance operation, updating affine parameters, but not updating running_stats. I think calling bn.eval() will utilize running_stats in forward pass of bachNorm layer. Is it right?</p>

ptrblck: <p>In that case you could use <code>F.linear</code> with the <code>bn.weight</code> and <code>bn.bias</code> during the forward pass.</p>

Mirsadeghi: <p>You mean something like:</p>
<aside class="quote no-group" data-username="Mirsadeghi" data-post="1" data-topic="104698">
<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/mirsadeghi/48/31772_2.png" class="avatar"> Mirsadeghi:</div>
<blockquote>
<p>x = extra_BatchNorm(x)*base_BatchNorm.weight+base_BatchNorm.bias</p>
</blockquote>
</aside>
<p>My issue is that, I want to apply this operation to other deep nets (resnet) which I can not directly change their forward pass methods as there are multiple use of bachNorm. Besides, I have to do this for every net and it is not possible. Can I apply this call:</p>
<aside class="quote no-group" data-username="Mirsadeghi" data-post="1" data-topic="104698">
<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/mirsadeghi/48/31772_2.png" class="avatar"> Mirsadeghi:</div>
<blockquote>
<p>x = extra_BatchNorm(x)*base_BatchNorm.weight+base_BatchNorm.bias</p>
</blockquote>
</aside>
<p>for bachnorm layers without changing their code internally. Something like ‘model.apply’?</p>

ptrblck: <p>No, <code>model.apply</code> will recursively apply the passed method to all layers.</p>
<p>I think the cleanest way would be to write a custom module using your suggested behavior and replace all batchnorm layers in the model with your custom ones.</p>

Mirsadeghi: <p>Sounds good. Thank you.</p>

swilson314: <p>I’m trying to freeze the bn training at a certain point, but I was surprised to find the affine parameters are updated when bn.eval(), this is by design?</p>

ptrblck: <p>Yes, as explained here:</p><aside class="quote" data-post="2" data-topic="104698">
  <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/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/affine-and-momentum-of-batchnorm-layer/104698/2">Affine and momentum of BatchNorm layer</a> 
  </div>
  <blockquote>
    You can call bn.eval() which will use the running stats and will not update them. The affine parameters will still be trained.
  </blockquote>
</aside>

<p>To freeze the affine parameters (i.e. the <code>weight</code> and <code>bias</code>) you would need to call <code>.requires_grad = False</code> on them.</p>

swilson314: <p>Thanks for your quick response. I certainly don’t see that in the documentation <a href="https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html" class="inline-onebox" rel="noopener nofollow ugc">BatchNorm2d — PyTorch 1.13 documentation</a>. So even during the my test() loop where I’ve set model.eval() the affine parameters are still being updated?!?</p>

ptrblck: <p>Yes, as already mentioned.<br>
You are right that the docs don’t seem to explain it in detail. Would you be interested in improving these?</p>

swilson314: <p>Sorry, to nag, but why would any parameters be updated during inference? Any why not documented?</p>

ptrblck: <p>Because calling <code>eval()</code> on any module is irrelevant to Autograd and only changed the behavior of some layers. This is by design to e.g. allow dropout layers to still be active during inference, which might be useful for some use cases. The affine, trainable parameters use the <code>.requires_grad</code> attribute to define if gradients should be computed for them or not (i.e. if these should be frozen or not).</p>
<aside class="quote no-group" data-username="swilson314" data-post="12" data-topic="104698">
<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/swilson314/48/29724_2.png" class="avatar"> swilson314:</div>
<blockquote>
<p>Any why not documented?</p>
</blockquote>
</aside>
<p>Improvements are more than welcome. Let me know if you are interested in improving the documentation.</p>

swilson314: <p>Ah, I see some light at the end of the tunnel!! Yes, I’d be happy to help with the documentation.</p>
<p>Can you help me understand a bit more? I guess I thought the actual parameter updating happened in response to backward and the optimizer.</p>
<p>I didn’t realize that dropout still worked during inference. I’m surprised there isn’t something like ireallyreallywanttoeval(true) that would turn off everything. I suspect that at the point that I get to an onnx model that dropout and bn are fully turned off?</p>

ptrblck: <p>I think you misunderstood me or I wasn’t clear enough.<br>
The <code>train()</code> and <code>eval()</code> calls on a model or module will set the internal <code>self.training</code> flag to <code>True</code> or <code>False</code>, respectively.<br>
This attribute changes the layer “behavior”, e.g. calling <code>dropout.eval()</code> will disable it, calling <code>batchnorm.eval()</code> will use the running stats to normalize the input activation and will not update the running stats anymore.</p>
<p>These calls, however, do not change the Autograd behavior and to not freeze trainable <code>nn.Parameter</code>s. These parameters, <code>weight</code> and <code>bias</code> in batchnorm layers, are frozen by changing their <code>requires_grad</code> attribute.</p>
<p>I tried to explain why the Autograd behavior is decoupled from the <code>train/eval</code> behavior, since you might want to use some layers in training behavior (e.g. enabling dropout), others in eval (e.g. in batchnorm layers) while freezing some and training other parameters.</p>

swilson314: <p>Thanks so much for the clarification. That’s in line with my understanding.</p>
<p>The comment I think may not be correct is this: This is by design to e.g. allow dropout layers to still be active during inference, which might be useful for some use cases. Per the documentation (<a href="https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html" class="inline-onebox" rel="noopener nofollow ugc">Dropout — PyTorch 1.13 documentation</a>) dropout only occurs during training, there is nothing to suggest it occurs during inference. Sorry to be pedantic, just want to make sure I understand.</p>

swilson314: <p>Maybe what would help the documentation is an explicit callout for what train() does.</p>

ptrblck: <aside class="quote no-group" data-username="swilson314" data-post="16" data-topic="104698">
<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/swilson314/48/29724_2.png" class="avatar"> swilson314:</div>
<blockquote>
<p>The comment I think may not be correct is this: This is by design to e.g. allow dropout layers to still be active during inference, which might be useful for some use cases. Per the documentation (<a href="https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html">Dropout — PyTorch 1.13 documentation</a>) dropout only occurs during training, there is nothing to suggest it occurs during inference. Sorry to be pedantic, just want to make sure I understand.</p>
</blockquote>
</aside>
<p>This is what I tried to clarify:</p>
<aside class="quote no-group" data-username="ptrblck" data-post="15" data-topic="104698">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>I think you misunderstood me or I wasn’t clear enough.<br>
The <code>train()</code> and <code>eval()</code> calls on a model or module will set the internal <code>self.training</code> flag to <code>True</code> or <code>False</code>, respectively.<br>
This attribute changes the layer “behavior”, e.g. calling <code>dropout.eval()</code> will disable it, calling <code>batchnorm.eval()</code> will use the running stats to normalize the input activation and will not update the running stats anymore.</p>
</blockquote>
</aside>
<p>I don’t claim dropout is always on, as it depends on the <code>self.training</code> status. Just that Autograd is completely independent from <code>train/eval</code> mode and you can mix different behavior together depending on your use case.<br>
If we forced to disable Autograd <em>and</em> forced to call <code>model.eval()</code> on all submodules, it would limit specific use cases such as explicitly allowing dropout to be used during testing.</p>

swilson314: <p>Got it! Thanks so much for the clarification.</p>