Fine tune the last layer while fix other layers

JonOnEarth: <p>Hi, I have some problems with fine-tuning the last layer of a Neural network.<br>
Code is as followed:</p>
<pre><code class="lang-auto">for layer in self.layers[0:-1]: 
        for param in self.model.fc_layers[layer].parameters():
            param.requires_grad = False
        
        y_prime = self.model(X)
        loss = self.model.criterion(y_prime, y)
        
        self.model.optimize.zero_grad()
        self.model.zero_grad()
        
        loss.backward(retain_graph=1)
        self.model.optimize.step()
</code></pre>
<p>Model is a Neural network.<br>
It didn’t work as I fine-tune the last layer(re-train), I get the same result while I didn’t fix the parameters of any layers.</p>
<p>something is wrong here, but I can’t figure it out.</p>

ptrblck: <p>Could you check, if the parameters of the frozen layers get valid gradients after the <code>backward</code> call via:</p>
<pre><code class="lang-python">loss.backward()
print(self.model.frozen_layer.weight.grad)
print(self.model.last_layer.weight.grad)
</code></pre>
<p>If the frozen layers yield <code>None</code>, while the last linear layer yields a valid gradient, your code is working as expected.</p>

JonOnEarth: <p>hi, thx.<br>
It has the error: the model has no ‘‘frozen_layer’’ module.<br>
So I change it to</p>
<pre><code class="lang-auto">print(self.model.fc_layers[-2].weight.grad)
print(self.model.last_layer.weight.grad)
</code></pre>
<p>Since the second last layer is frozen as I wrote the code in question.<br>
But the output is not “None”, is the gradient matrix.<br>
So it is not working as expected I think.</p>

ptrblck: <p>Yes, the <code>frozen_layer</code> and <code>last_layer</code> are just placeholder names for your layer names.</p>
<p>Did you train the “frozen” layers before and forgot to zero out the gradients?<br>
If not, could you post a code snippet which represents your training routine and how and where you are freezing the layers?</p>
<p>Your initial code snippet looks like it has some indentation issues, as your “layer loop” would be used for each training iteration.<br>
Is this a copy-paste error or are you using this code in your script as shown here?</p>

JonOnEarth: <p>oh, sorry, yes, it’s a copy-paste mistake.</p>
<p>I train the frozen layer before, the function is as followed:</p>
<pre><code class="lang-auto">def step_mlp(self,X,y):
        y_prime = self(X)
        loss = self.criterion(y_prime, y)
        
        self.optimize.zero_grad()
        self.zero_grad()
        
        loss.backward(retain_graph=1)
        self.optimize.step()
</code></pre>
<p>Then re-train the model by frozen the layers except the last layer as the code in question decription.<br>
Frozen layer:</p>
<pre><code class="lang-auto">for layer in self.layers[0:-1]: 
        for param in self.model.fc_layers[layer].parameters():
            param.requires_grad = False
</code></pre>
<p>Then re-train:</p>
<pre><code class="lang-auto">y_prime = self.model(X)
loss = self.model.criterion(y_prime, y)
        
self.model.optimize.zero_grad()
self.model.zero_grad()
loss.backward(retain_graph=1)
self.model.optimize.step()
</code></pre>

ptrblck: <p>Your general workflow should work as shown in this small example:</p>
<pre><code class="lang-python"># Setup
model = models.resnet18()
data = torch.randn(1, 3, 224, 224)
target = torch.randint(0, 1000, (1,))
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# Train
optimizer.zero_grad()
out = model(data)
loss = criterion(out, target)
loss.backward()
optimizer.step()

# Freeze all but last layer
for name, param in model.named_parameters():
    if not 'fc' in name:
        param.requires_grad = False
      
optimizer.zero_grad()
out = model(data)
loss = criterion(out, target)
loss.backward()

# Check grads
grad_frozen = model.conv1.weight.grad
grad_fc = model.fc.weight.grad
print(grad_frozen.abs().sum())
print(grad_fc.abs().sum())

optimizer.step()
</code></pre>
<p>As you can see, after freezing all parameters but the <code>fc.weight</code> and <code>fc.bias</code>, I get a zero gradient for a previous layer, while I get valid gradients for the last linear layer.</p>
<p>Which layers are you checking and are you sure the checked frozen layer is in <code>model.fc_layers</code>?</p>

JonOnEarth: <p>Thanks for your detailed example, ptrblck!</p>
<p>Find the error. It turns out that I didn’t have any frozen layers in <code>model.fc_layers</code> because I make a mistake in my code.<br>
Your reply helps me to find it.  Thanks a lot!</p>