How to freeze BN layers while training the rest of network (mean and var wont freeze)

seyeeet: <p>I have a network that consists of batch normalization (BN) layers and other layers (convolution, FC, dropout, etc)<br>
I was wondering how we can do the following :</p>
<ol>
<li>I want to freeze all the layer and just train the BN layers</li>
<li>freeze the BN layers and train every other layer in the network except BN layers</li>
</ol>
<p>My main issue is how to handle freezing and training the BN layers</p>

Usama_Hasan: <p>You can do something like this, hope this works for you.</p>
<pre><code class="lang-auto">model = Net()
for name ,child in (model.named_children()):
    if name.find('BatchNorm') != -1:
        for param in child.parameters():
            param.requires_grad = True
    else:
        for param in child.parameters():
            param.requires_grad = False 
</code></pre>

seyeeet: <p>thank you for your reply, but the <code>bn.running_mean</code>  and  <code>bn.running_var</code> values of the BN still change when I freeze them</p>

ptrblck: <p>The running stats are not parameters (they do not get gradients and the optimizer won’t update them), but buffers, which will be updated during the forward pass.<br>
You could call <code>.eval()</code> on the batchnorm layers to apply the running stats instead of calculating and using the batch statistics and updating the running stats (which would be the case during <code>.train()</code>).</p>

seyeeet: <p>thanks for your reply. can you please give me a simple snippet code on how I can train the network while freezing the BN layer?</p>

ptrblck: <p><a class="mention" href="/u/usama_hasan">@Usama_Hasan</a>’s code snippet could be used to freeze the affine parameters of all batchnorm layers, but I would rather use <code>isinstance(child, nn.BatchNorm2d)</code> instead of the name check.<br>
Anyway, his code should also work. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Usama_Hasan: <p>Yes, <a class="mention" href="/u/ptrblck">@ptrblck</a> thanks for the advice, it’s  a more pythonic and stable solution.</p>

seyeeet: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="89736">
<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>.eval()</p>
</blockquote>
</aside>
<p>usually, I use <code>model.train()</code> when training, but now I don’t wanna train the BN during the training. does it sounds right if in training I do  <code>model.eval()</code> and then do the following</p>
<pre><code class="lang-auto">
for name ,child in (model.named_children()):
    if name.find('BatchNorm') != -1:
        for param in child.parameters():
            param.requires_grad = False
    else:
        for param in child.parameters():
            param.requires_grad = True
</code></pre>

ptrblck: <p>The code looks alright, but again I would recommend to use <code>isinstance</code> instead of the name check. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"><br>
Note that <code>model.eval()</code> will also disable dropout layers (any maybe change the behavior of custom layers), so you might want to call <code>eval()</code> only on the batchnorm layers.</p>

seyeeet: <p>oh I see,<br>
Didn’t know I can call eval() just on BN layers<br>
Will give it a try with eval on BN layers and will use <code>isinstance</code><br>
thanks for clarifying</p>

isalirezag: <p>here is what you want to do.</p>
<p>You dont need to use both hasattr and eval together. But I just did it to be safe.</p>
<pre><code class="lang-auto">for module in model.modules():
    # print(module)
    if isinstance(module, nn.BatchNorm2d):
        if hasattr(module, 'weight'):
            module.weight.requires_grad_(False)
        if hasattr(module, 'bias'):
            module.bias.requires_grad_(False)
        module.eval()
</code></pre>

Ahmed_m: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I use module.eval() to freeze the stats of an instancenorm layer (with track_running_stats = True) after some epochs. However, the running_mean and ranning_var are still updating over training. Any idea why this may happen?</p>

ptrblck: <p>I cannot reproduce this issue using this code snippet and the running stats are equal after calling <code>norm.eval()</code>:</p>
<pre><code class="lang-python">norm = nn.InstanceNorm2d(num_features=3, track_running_stats=True)
print(norm.running_mean, norm.running_var)
&gt; tensor([0., 0., 0.]) tensor([1., 1., 1.])

x = torch.randn(2, 3, 24, 24)

out = norm(x)
print(norm.running_mean, norm.running_var)
tensor([-0.0029,  0.0005,  0.0003]) tensor([0.9988, 1.0021, 0.9980])

out = norm(x)
print(norm.running_mean, norm.running_var)
&gt; tensor([-0.0056,  0.0010,  0.0006]) tensor([0.9978, 1.0040, 0.9962])

norm.eval()
out = norm(x)
print(norm.running_mean, norm.running_var)
&gt; tensor([-0.0056,  0.0010,  0.0006]) tensor([0.9978, 1.0040, 0.9962])
</code></pre>
<p>Are you using an older PyTorch version, where this could have been a known issue?</p>

Ahmed_m: <p>I use version 1.6 and the issue arise when I switch to eval() during the training. I have a situation like the following:</p>
<pre><code class="lang-auto">def _freeze_norm_stats(net):
    try:
        for m in net.modules():
           
            if isinstance(m, InstanceNorm1d):
                
                m.track_running_stats = False
                m.eval()

    except ValueError:  
        print("errrrrrrrrrrrrrroooooooorrrrrrrrrrrr with instancenorm")
        return

model = Model().cuda()

for epoch in range(start_epoch,end epoch):
    #for epochs &lt; 10, the model updates the running_mean and running_var of its instancenorm layer
    if epoch == 10:
        print("freeze stats:")
        model.apply(_freeze_norm_stats)
    #for epochs &gt;= 10, the model should continue with fixed fixed mean and var for its instancenorm layer

</code></pre>
<p>Unfortunately, printing the running_mean and running_mean after calling the function _freeze_norm_stats shows that they’re still being updated over training steps. Maybe I have something wrong in the function implementation but I am still not able to figure this out.</p>

ptrblck: <p>Could you update to the latest stable release (<code>1.7.1</code>) and check, if you are still facing this issue?</p>

Ahmed_m: <aside class="quote no-group" data-username="Ahmed_m" data-post="15" data-topic="89736">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/4491bb/48.png" class="avatar"> Ahmed_m:</div>
<blockquote>
<p><code>m.track_running_stats = False</code></p>
</blockquote>
</aside>
<p>removing this line and only keeping the .eval() solved the problem.</p>

ss007k: <p>What could be the easiest way to freeze the batchnorm layers in say, layer 4 in Resnet34? I am finetuning only layer4, so plan to check both with and without freezing BN layers.</p>
<p>I checked resnet34.layer4.named_children() and can write loops to fetch BN layers inside layer4 but want to check if there is a more elegant way.</p>

yclin: <p>Encounter the same issue: the running_mean/running_var of a batchnorm layer are still being updated even though “bn.eval()”. Turns out that the only way to freeze the running_mean/running_var is “bn.track_running_stats = False” . Tried 3 settings:</p>
<ul>
<li>bn.param.requires_grad = False &amp; bn.eval()</li>
<li>bn.param.requires_grad = False &amp; bn.track_running_stats = False &amp; bn.eval()</li>
<li>bn.param.requires_grad = False &amp; bn.track_running_stats = False</li>
</ul>
<p>The first one did not work. Both 2nd and 3rd worked. So my conclusion is that  use “bn.track_running_stats = False” if you wanna freeze running_mean/running_var. (Contradictory to <a href="https://discuss.pytorch.org/t/how-to-freeze-bn-layers-while-training-the-rest-of-network-mean-and-var-wont-freeze/89736/17" class="inline-onebox">How to freeze BN layers while training the rest of network (mean and var wont freeze) - #17 by Ahmed_m</a></p>
<p>Additional info: torch.<strong>version</strong>: 1.7.0. Seems there is a misalignment across versions.</p>

liyuecheng27: <p>I tried your way and it absolutely works with setting track_running_stats=False. However, just out of curiosity, in the official document from PyTorch, it said if track_running_stats is set to False, then running_mean and running_var will be none and the batch norm will always use batch stat?<br>
My guess is that track_running_stats will not reset running_mean and running_var, and hence what it only does is not update them anymore. Please, correct me if my assumption is wrong.</p>

ptrblck: <p>Yes, using <code>track_running_stats=False</code> will always normalize the input activation with the current batch stats so you would have to make sure that enough samples are passed to the model (even during testing).</p>