Freezing the parts of the model by passing only certain parameters when creating optimizers?

username: Kore_ana
cooked: <p>Hi! I’m trying to freeze parts of the code and I tried the method of only putting in certain parameters when  defining the optimizer, which seems to work but I couldn’t find this method anywhere else on the web so was not sure if this method is OK to use… the code is below. any help would be greatly appreciated!</p>
<pre><code class="lang-auto">class mymodel(nn.Module):
    def __init__(self):
        super().__init__() 
        self.layer1 = nn.Linear(10, 5)
        self.layer2 = nn.Linear(5,5)
        self.layer3 = nn.Linear(5,5)
        self.act = nn.Sigmoid()
    def forward(self, x):
        x = self.layer1(x) 
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.act(x)
        return x
model = mymodel()
</code></pre>
<pre><code class="lang-auto">optimizer = optim.Adam(model.parameters(), lr=0.001)
</code></pre>
<p>when defining like this, i find that layer 1,2,3 have different weight values, while doing like :</p>
<pre><code class="lang-auto">finetune_1_optimizer = optim.Adam(model.layer2.parameters(), lr=0.001)

</code></pre>
<p>only changes layer2’s weight values. (albeit the “requires_grad” still being true)</p>
<p>Is this method of freezing parts of the model OK? I though this method wouldn’t work because if I (for example) only put in layer 2 into the optimizer, since the computation graph is still linked, when backprop occurs the grad_fn values in layer 1~3 will all be affected, meaning that all of them could change.</p>

username: KFrank
cooked: <p>Hi Kore!</p>
<aside class="quote no-group" data-username="Kore_ana" data-post="1" data-topic="167082" data-full="true">
<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/kore_ana/48/45740_2.png" class="avatar"> Kore_ana:</div>
<blockquote>
<pre><code class="lang-auto">finetune_1_optimizer = optim.Adam(model.layer2.parameters(), lr=0.001)

</code></pre>
<p>only changes layer2’s weight values. (albeit the “requires_grad” still being true)</p>
<p>Is this method of freezing parts of the model OK?</p>
</blockquote>
</aside>
<p>Yes, this is a sound way to freeze the weights of the other layers.</p>
<aside class="quote no-group">
<blockquote>
<p>I though this method wouldn’t work because if I (for example) only put in layer 2 into the optimizer, since the computation graph is still linked, when backprop occurs the grad_fn values in layer 1~3 will all be affected,</p>
</blockquote>
</aside>
<p>It is true that calling <code>.backward()</code> on you loss will cause the gradients of<br>
all of the layers to be populated.</p>
<aside class="quote no-group">
<blockquote>
<p>meaning that all of them could change.</p>
</blockquote>
</aside>
<p>But populating the gradients doesn’t, in and of itselt, cause the weights<br>
to change.  It’s the optimizer that changes the weights and because<br>
<code>finetune_1_optimizer</code> has only had <code>layer2</code>’s weights added to it,<br>
calling <code>finetune_1_optimizer.step()</code> will only modify <code>layer2</code>’s weights<br>
and will leave <code>layer1</code> and <code>layer3</code> unchanged.</p>
<p>Note, this approach has two consequences:</p>
<p>You will pay the price – which may or may not matter – of backpropagating<br>
through <code>layer1</code> and updating its gradients.  (You do need to backpropagate<br>
through <code>layer3</code> in order to compute the gradients for <code>layer2</code>.)</p>
<p><code>loss.backward()</code> <em>accumulates</em> the gradients of all three layers.  if you<br>
call <code>finetune_1_optimizer.zero_grad()</code>, it will only set the gradients<br>
for <code>layer2</code>, so you could repeatedly be accumulating into the gradients<br>
for the other two layers without ever zeroing them out.  This is unlikely to be<br>
an issue, but it could lead to overflow (that shouldn’t matter, anyway).  You<br>
could call <code>model.zero_grad()</code> or <code>optimizer.zero_grad()</code> to zero out<br>
all three layers’ weights.</p>
<p>As an aside, the way your example <code>mymodel</code> is written – without any<br>
intervening nonlinearities between your three <code>Linear</code> layers – those<br>
three layers collapse, in effect, into a single <code>Linear (10, 5)</code> layer.  In<br>
a real model you would want some sort of activation, <code>ReLU</code> or <code>Sigmoid</code>,<br>
say, in between the <code>Linear</code> layers.</p>
<p>Best.</p>
<p>K Frank</p>

username: Kore_ana
cooked: <p>Thank you for your answer <a class="mention" href="/u/kfrank">@KFrank</a> ! I have a further question if you don’t mind : I found that when freezing the model as my original question, the (supposed to be frozen) variables like <code>running_var</code> of the <code>BatchNorm</code> (which I set as frozen) also changed, which seems like it could be problematic as I do not want the parameters of the BatchNorm to change during evaluation. In this case, should I change the batch norm layers to eval mode? (I was going to optimize the code by only leaving the unfrozen layers to <code>.train()</code> and the rest to <code>.eval()</code> anyway, but I’m just curious if my understanding is correct! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> )</p>