Is it necessary to set BatchNorm to eval mode when we finetune the last layer of a pretrained model?

SunHaozhe: <p>Suppose that I want to use a pretrained <a href="https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html" rel="noopener nofollow ugc">ResNet18</a> to do classification on a small dataset, I just want to train the last (linear) classification layer while freezing all previous layers (referred to as “freezed” layers). This is a fairly common use case. To do so, the <a href="https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html#convnet-as-fixed-feature-extractor" rel="noopener nofollow ugc">official PyTorch tutorial</a> says that it is sufficient to do:</p>
<pre><code class="lang-auto">model_conv = torchvision.models.resnet18(pretrained=True)
for param in model_conv.parameters():
    param.requires_grad = False

num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)

...

# training phase
model.train()
for inputs, labels in dataloaders:
    optimizer.zero_grad()
    outputs = model(inputs)
    ...
    loss.backward()
    optimizer.step()
    ...
</code></pre>
<p>The tutorial says that it is sufficient to set <code>requires_grad</code> to <code>False</code> for “freezed” layers, it did not call <code>.eval()</code> for <code>BatchNorm2d</code> of “freezed” layers of resnet18. Does this mean that even if convolutional layers and affine parameters of BatchNorm2d of “freezed” layers are not learning, the running average of BatchNorm in “freezed” layers are still being updated even if we only wanted to train the last linear layer? Is this a mistake?</p>
<p>According to the following posts and documentation, it seems that in addition to set <code>requires_grad</code> to <code>False</code> for “freezed” layers (convolutional layers and BatchNorm layers), we should also call <code>.eval()</code> on all BatchNorm layers if we only want to train the last linear layer while freezing all “freezed” layers, which is contradicting the <a href="https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html#convnet-as-fixed-feature-extractor" rel="noopener nofollow ugc">official PyTorch tutorial on finetuing</a>.</p>
<pre><code class="lang-auto">def set_bn_eval(module):
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
        module.eval()
        
model.apply(set_bn_eval)
</code></pre>
<ul>
<li><a href="https://discuss.pytorch.org/t/should-i-use-model-eval-when-i-freeze-batchnorm-layers-to-finetune/39495">Should I use model.eval() when I freeze BatchNorm layers to finetune?</a></li>
<li><a href="https://discuss.pytorch.org/t/finetuning-pretrained-networks-with-batchnorm/4121">Finetuning pretrained networks with batchnorm</a></li>
<li><a href="https://discuss.pytorch.org/t/how-to-fix-batchnorm-parameters-during-finetune/57304">How to fix BatchNorm parameters during finetune</a></li>
<li><a href="https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html" rel="noopener nofollow ugc">torch.nn.BatchNorm2d documentation</a></li>
</ul>
<p>Do we indeed need to call <code>.eval()</code> on BatchNorm layers when we only train the last linear layer of a pretrained resnet18 model?</p>

ptrblck: <aside class="quote no-group" data-username="SunHaozhe" data-post="1" data-topic="156031">
<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/sunhaozhe/48/14036_2.png" class="avatar"> SunHaozhe:</div>
<blockquote>
<p>Does this mean that even if convolutional layers and affine parameters of BatchNorm2d of “freezed” layers are not learning, the running average of BatchNorm in “freezed” layers are still being updated even if we only wanted to train the last linear layer?</p>
</blockquote>
</aside>
<p>Yes, the running stats would still be update in training mode.</p>
<aside class="quote no-group" data-username="SunHaozhe" data-post="1" data-topic="156031">
<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/sunhaozhe/48/14036_2.png" class="avatar"> SunHaozhe:</div>
<blockquote>
<p>Is this a mistake?</p>
</blockquote>
</aside>
<p>I don’t know if it’s a mistake and would guess it could help in fine tuning the model especially if your new input data has different stats.</p>