Should I use model.eval() when I freeze BatchNorm layers to finetune?

fanl: <p>Hi,<br>
I have a well trained coarse net (including BN layers) which I want to freeze to finetune other layers added.  I filtered out the parameters of the coarse net when construct optimizer.  Is model.eval() or something else necessary in this case？I don’t want the BN layers to recalculate the mean and variance in every batch.</p>

ptrblck: <p>If you set your <code>nn.BatchNorm</code> layers to <code>eval()</code> the running estimates won’t be updated anymore.<br>
Additionally to filtering out the parameters, you could also set the <code>.requires_grad</code> attribute to <code>False</code>, so that the gradients won’t be computed if not necessary.</p>

fanl: <p>Thanks!!  But I am still a little confused.</p>
<pre><code class="lang-auto">for p in model.parameters():
    p.requires_grad = False
for p in model.fine.parameters():
    p.requires_grad = True
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),lr = 0.001)

</code></pre>
<p>This is my code for finetuning, and I didn’t use model.eval().  Do the BN layers in my model behave the same as BN layers in other models whose eval() are true?<br>
If it’s not same, what should I do to freeze the BN layers (make BN layers use global means and variances instead of them of every mini-batch)?<br>
I am looking forward to your reply.</p>

imaluengo: <p>If you only want to fine-tune the parameters in <code>model.fine</code> you could do the following:</p>
<pre><code class="lang-auto">model = ...
optimizer = torch.optim.Adam(model.fine.parameters(),lr = 0.001)
</code></pre>
<p>Now the optimizer will only try to update parameters within the <code>.fine</code> submodule.</p>

ptrblck: <p>It depends if they were set to <code>.eval()</code> before, but the default mode is <code>train()</code> after loading the model.<br>
If you want to set the complete model to <code>eval</code> mode, just use <code>model.eval()</code>.<br>
Alternatively, if you just want to apply it on all batch norm layers, you could use:</p>
<pre><code class="lang-python">def set_bn_eval(module):
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
        module.eval()
        
model.apply(set_bn_eval)
</code></pre>

Rakshit_Kothari: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> a little more detail on this - does setting the batch norm layer to <code>eval</code> allow us to train ‘gamma’ and <code>beta</code> parameters? I understand that the <code>eval</code> operation allows us to use the current batch’s mean and variance when fine tuning a pretrained model.</p>

ptrblck: <p>The <code>.train()</code> and <code>.eval()</code> call on batchnorm layers does not freeze the affine parameters, so that the gamma (<code>weight</code>) and beta (<code>bias</code>) parameters can still be trained.</p>
<aside class="quote no-group" data-username="Rakshit_Kothari" data-post="6" data-topic="39495">
<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/rakshit_kothari/48/1974_2.png" class="avatar"> Rakshit_Kothari:</div>
<blockquote>
<p>I understand that the <code>eval</code> operation allows us to use the current batch’s mean and variance when fine tuning a pretrained model.</p>
</blockquote>
</aside>
<p>Calling <code>eval()</code> on batchnorm layers will use the running stats, while <code>train()</code> will use the batch stats and update the running stats.</p>

Rakshit_Kothari: <p>Thanks, that clarification was useful.</p>

deeptigp: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> - one follow up on your response earlier. Does setting <code>requires_grad = False</code> sufficient, even for batch norm layers? Or do we have to do both, i.e., set <code>requires_grad = False</code> AND <code>BN_layer.eval()</code>?</p>

ptrblck: <p>The <code>requires_grad</code> attribute and calling <code>train()/eval()</code> on it behave differently.<br>
BatchNorm layers use trainable affine parameters by default, which are assigned to the <code>.weight</code> and <code>.bias</code> attribute. These parameters use <code>.requires_grad = True</code> by default and you can freeze them by setting this attribute to <code>False</code>.<br>
During training (i.e. after calling <code>model.train()</code> or after creating the model) the batchnorm layer will normalize the input activation using the batch stats and will update the internal stats using a running average.<br>
After calling <code>model.eval()</code> the batchnorm layers will use the trained internal running stats (stored as <code>.running_mean</code> and <code>.running_var</code>) to normalize the input activation.</p>

deeptigp: <p>Got it, thanks so much for your detailed response! So, in the event that I set <strong>just</strong> <code>requires_grad = False</code> for BN layers, it may still be computing the running average during training phase, which is not ideal. I should be doing both, so the batchnorm layer uses the stored <code>.running_mean</code> and <code>.running_var</code> values for normalization.</p>

ptrblck: <p>If you want to validate your model, wrapping the forward pass into <code>with torch.no_grad()</code> or <code>with torch.inference_mode()</code> and calling <code>model.eval()</code> would also work. You wouldn’t necessarily need to flip the <code>.requires_grad</code> attribute (it would also work, but the former guards might be a simpler way).</p>

jia_lee: <p>If I want to fine-tune a model with BN layers and only freeze the BN layers, we should set <code>model.eval()</code> and BN layers’ requires_grad to False. Am I right?</p>

ptrblck: <p>Yes, if you want to freeze the affine parameters (<code>weight</code> and <code>bias</code>), set their <code>.requires_grad</code> attribute to <code>False</code>). If you additionally want to use the running stats and not update them with the current batch stats, call <code>.eval()</code> on the batchnorm layers.</p>