Problems with the optimiser when resuming

MauroGentile: <p>I have a NN with 4 heads that I want to train one after the other.</p>
<p>After having trained 1 head, I save the net and I then want to resume the model, change the requires_grad of the already trained head and start training the next head but I get this error:<br>
ValueError: loaded state dict contains a parameter group that doesn’t match the size of optimizer’s group</p>
<p>I get the same error even if  the requires_grad is changed after having loaded the model.</p>
<p>I do not get any error if I do not make any change to requires_grad. Meaning that resume works fine.</p>
<p>The problem arises when I load the optimiser. The model can be loaded without problems.</p>
<p>Observe that the the Net structure never changes in time. It is the trainability of some layers that are switched.</p>
<p>How can I solve this problem?</p>

ptrblck: <p>Are you filtering the parameters based on their <code>requires_grad</code> attribute before passing them to the optimizer? Also, could you post an executable code snippet to reproduce this issue, as I’m currently unsure what might cause the error?</p>

MauroGentile: <p>Hello, first of all thanks for helping.</p>
<p>My code is structured in this way:</p>
<ol>
<li>I initially create the net structure, with all param_grad=True by default</li>
<li>depending on information on a config file I switch to False the param_grad of all heads but one. For instance the code to switch to False the heat map head is the following:</li>
</ol>
<pre><code class="lang-auto">if config["trainable_heads"]["heatmap"] == False:
      enc_hm = map(lambda x: x[1],
                                  filter(lambda p: p[1].requires_grad and not ("backbone.decoder_delta" in p[0]),
                                         model.named_parameters()))
      for param in enc_hm:
          param.requires_grad = False 
</code></pre>
<p>I have 3 other identical snippet for the other 3 heads.</p>
<p>Then I train and save the net.<br>
When I resume, I get the error, which only happens if I try to switch the trainable head. If I do not switch heads, the resume works perfectly.</p>
<p>An attempt I made is the following: before resuming from previous training, I set the layers of the network in the exact same  trainability status as just before saving. I checked it was the case by comparing the name of all trainable parameters just before saving and just before resuming. The 2 list are indeed identical and as I would have expected. Which confuses me a lot.</p>
<p>The error is when I load the optimiser:</p>
<pre><code class="lang-auto">            self.model.load_state_dict(checkpoint['state_dict'])
            self.scheduler.load_state_dict(checkpoint['scheduler']) 
            self.optimizer.load_state_dict(checkpoint['optimizer'])
</code></pre>
<p>No error is raised in the first 2.</p>

ptrblck: <p>I cannot reproduce the issue using different work flows of freezing different sets of parameters:</p>
<pre><code class="lang-python">model = models.resnet18()
for param in model.fc.parameters():
    param.requires_grad = False

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

output = model(torch.randn(1, 3, 224, 224))
output.mean().backward()
optimizer.step()
optimizer.zero_grad()

torch.save(model.state_dict(), 'model.pt')
torch.save(optimizer.state_dict(), 'opt.pt')

# load plain model and optimizer
model = models.resnet18()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model.load_state_dict(torch.load('model.pt'))
optimizer.load_state_dict(torch.load('opt.pt'))

# load model with same frozen parameters
model = models.resnet18()
for param in model.fc.parameters():
    param.requires_grad = False
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model.load_state_dict(torch.load('model.pt'))
optimizer.load_state_dict(torch.load('opt.pt'))

# load model with different frozen parameters
model = models.resnet18()
for param in model.conv1.parameters():
    param.requires_grad = False
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model.load_state_dict(torch.load('model.pt'))
optimizer.load_state_dict(torch.load('opt.pt'))
</code></pre>

MauroGentile: <p>Thank you. I will go through my code once again. There must be some bug there.</p>

MauroGentile: <p>Hello,</p>
<p>From what I see from your code, you always pass all networks parameters to the Adam optimizer, regardless requires_grad value of each of them.</p>
<p>I thought I had to pass to the optimiser only the trainable ones. This is incorrect, right?</p>
<p>Thanks for helping</p>

ptrblck: <p>You could filter them out before, if you like. As long as the <code>.grad</code> attributes aren’t filled in the frozen parameters, the optimizer won’t update them, but your approach would be more explicit.<br>
Were you able to reproduce the issue using my code snippet and the filtering or any other addition to the code?</p>