How to freeze BN layers while training the pretrained model

Yangmin: <p>I have a network that consists of batch normalization (BN) layers and other layers (convolution, FC, dropout, etc) which is pretrained ResNet50 model.</p>
<p>I want the model not to be trained so I freezed the all layer with requires_grad=False,<br>
but I find the BN layer still updating and the performance gradually dropped.</p>
<p>So, I used the below code to freeze the batch norm layer.</p>
<pre><code>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.track_running_stats = False
        # module.eval()
</code></pre>
<ol>
<li>
<p>But I am confused about the difference between putting layer.eval() and module.track_running_stats=False.</p>
</li>
<li>
<p>If I want the BN layer not to be updated and performance not to be changed in inference, do I just need to make it module.weight.requires_grad(False) only? Or should I also stop track_running_stats<br>
which is used in inference mode.</p>
</li>
</ol>
<p>Thank you.</p>

Yangmin: <p>I just figured out if I used the code below, which do not put layer into .eval() but sets track_running_stats=False only,<br>
I get an error.</p>
<p>for module in model.modules():</p>
<pre><code>    # 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.track_running_stats = False
        #module.eval()</code></pre>

ptrblck: <p>I don’t know what kind of error you are getting, but changing the <code>track_running_stats</code> attribute after the layer creation might be dangerous. It is working in my setup, but I don’t think there is a guarantee that changing this attribute would work in previous/next versions.</p>
<pre><code class="lang-python">bn = nn.BatchNorm2d(3)
print(bn.running_mean)
&gt; tensor([0., 0., 0.])

out = bn(torch.randn(2, 3, 24,  24))
print(bn.running_mean)
&gt; tensor([0.0056, 0.0054, 0.0008])

bn.track_running_stats = False
out = bn(torch.randn(2, 3, 24,  24))
print(bn.running_mean)
&gt; tensor([0.0056, 0.0054, 0.0008])
</code></pre>
<p>The proper approach would be to use <code>module.eval()</code> or is this not working for you?</p>

Yangmin: <p>Putting module.eval() works and does not change the performance.<br>
But the bn layer keep changning running_mean and running_var during training even though I freezed them with module.eval().</p>
<p>So I thought I should also turn the track_running_stats off to prevent the running_mean and running_var which is used in inference.</p>

ptrblck: <aside class="quote no-group" data-username="Yangmin" data-post="4" data-topic="135630">
<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/yangmin/48/43502_2.png" class="avatar"> Yangmin:</div>
<blockquote>
<p>But the bn layer keep changning running_mean and running_var during training even though I freezed them with module.eval().</p>
</blockquote>
</aside>
<p>That shouldn’t be the case. Could you post a minimal, executable code snippet showing this behavior, please?</p>

Yangmin: <p>You are right. I thought the running_mean and running_var was changing<br>
because the performance dropped.</p>
<p>When setting BN layer module.eval() the running_mean and running_var does not change.</p>