Optimizer.load_state_dict() weird behaviour with Adam optimizer

Ajinkya_Ambatwar: <p>Hi all, I am trying to implement a recent ICCV paper in PyTorch. I have to train an enormous sized dataset(in 10s of GB) for very large number of epochs(4000). So I am trying to store the model and optimizer state after certain epochs to save the state and load it afterwards.</p>
<pre><code class="lang-auto">torch.save({
        'epoch': epoch,
        'model_state_dict': net.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss
        }, path)
</code></pre>
<p>Then loading it like this -</p>
<pre><code class="lang-auto">net = model(*args)
net = net.to(device)
optimizer = optim.Adam(net.parameters(), lr = 8e-5)
if train_epoch != None:
    checkpoint = torch.load(path)
    net.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    train_epoch = checkpoint['epoch']
    loss = checkpoint['loss']
</code></pre>
<p>I am observing that the loss from the loaded model if is 0.0756, when I start training with the loaded model, the loss after 1st iteration in the next epoch jumps to 0.23105.<br>
It would take some number of epochs at-least to come back to the same point. I am not sure what actually is causing the problem. Is it something related to the way Adam optimizer works based on previous gradient history.</p>
<p><strong>NOTE  -  I am zeroing the gradients in every epoch</strong> using</p>
<pre><code class="lang-auto"> optimizer.zero_grad()
</code></pre>
<p>Is that causing the error? What actually is the correct way to load the state dictionary for model and optimizer in case of any optimizer, Adam in particular?</p>
<p>Thanks!</p>

ptrblck: <p>Restoring the Adam optimizer should work as shown in this dummy example:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import torch.optim as optim

import torchvision.models as models


device = 'cuda'
model = models.resnet18()
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

data = torch.randn(100, 3, 224, 224, device=device)
target = torch.randint(0, 1000, (100, ), device=device)

nb_epochs = 10
for epoch in range(nb_epochs):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print('Epoch {}, loss {}'.format(epoch, loss.item()))

# Create reference loss
model.eval()
with torch.no_grad():
    output = model(data)
    ref_loss = criterion(output, target)
print('reference loss {}'.format(ref_loss))

# restore
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict()
}

model = models.resnet18()
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])

model.eval()

output = model(data)
loss = criterion(output, target)
print('restored loss {}'.format(loss.item()))
print('abs error {}'.format((ref_loss - loss).abs().item()))
</code></pre>
<p>Are you changing the train/eval mode or how do you compare your losses?</p>

Ajinkya_Ambatwar: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, thanks for the reply<br>
I think I am doing the same thing that you have mentioned. Is it something related to the way how ADAM works? Also at around 70th epoch I changed lr from 1e-4 to 8e-5 and trained further till 300 with the same lr. As you can see, I am not using any scheduler or weight_decay.<br>
For the loss, I have created a custom loss function looks like</p>
<pre><code class="lang-auto">Loss = torch.mean(torch.abs(out - target))
</code></pre>

ptrblck: <p>If you are just saving and restoring the <code>state_dicts</code> as in my example, it should work. Could you post your code so that we can have a look?</p>

Ajinkya_Ambatwar: <p>I am loading the 150th epoch checkpoint model, the loss is 0.0400 as can be seen.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/f/1f6669a8a16d5a31f841f18ce01bd7dbba434422.png" data-download-href="https://discuss.pytorch.org/uploads/default/1f6669a8a16d5a31f841f18ce01bd7dbba434422" title="Loss"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1f6669a8a16d5a31f841f18ce01bd7dbba434422_2_690x293.png" alt="Loss" data-base62-sha1="4tMdD57J508h2y5mcUvQyMSgXfQ" width="690" height="293" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1f6669a8a16d5a31f841f18ce01bd7dbba434422_2_690x293.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1f6669a8a16d5a31f841f18ce01bd7dbba434422_2_1035x439.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/f/1f6669a8a16d5a31f841f18ce01bd7dbba434422.png 2x" data-dominant-color="2E2D2D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Loss</span><span class="informations">1129×480 57.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Now I am running the main loop as follows -</p>
<pre><code class="lang-auto">for epoch in range(train_epoch+1, 50+train_epoch+1):
      for each data point in dataset:
       # Making the input and output ready from an image
        input_patch = (some random patch from input image)
        gt_patch = (same indexed patch from gt_image)
        optimizer.zero_grad()
        out = net(input_patch)
     
        loss = abs_L1_loss(out, gt_patch)
        loss.backward()
        optimizer.step()
        out = out.detach().cpu().numpy()
        out = np.minimum(np.maximum(out, 0), 1)
        # print("Out shape ="+str(out.shape))
        g_loss[ind] = loss.detach().cpu()

        print("%d %d Loss=%.5f Time=%.3f" % (epoch, cnt, np.mean(g_loss[np.where(g_loss)]), time.time() - st))

        
    torch.save({
        'epoch': epoch,
        'model_state_dict': net.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss
        }, checkpoint_dir+'/%04d-model.ckpt' %epoch)

</code></pre>
<p>The result captured is as follows -<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/5/0/50ee227ed09ad5408fef8d4f33a5425c68d6be37.png" alt="newloss" data-base62-sha1="bxWoJFZgyOGnGuCFMikHhFOPHyD" width="491" height="140"></p>
<p><strong>NOTE: I am taking a new random patch from the input image every iteration in every epoch</strong><br>
Can’t really understand why it is happening?</p>

Eric_K: <p>Hi ptrblck~</p>
<p>Sorry to bother you. I have a question to ask. What modification should I do to optimizer’s state_dict, if I modified the network’s structure?(e.g. I add several layers at the end of the network)</p>
<p>And what keys and values stand for in optimizer’s state_dict.(“state”, “param_groups”)</p>
<p>Thanks,<br>
Best,<br>
eric</p>

ptrblck: <p>If you modified the model before starting the training, you can just initialize a new optimizer as usual.<br>
On the other hand, if you modified it afterwards, you could add the parameters of the newly added modules via <code>optimizer.add_param_group</code>.</p>

HampusAstrom: <p>Was there ever a solution to the original posters problem (Ajinkya_Ambatwar)? I have a similar problem in my research project.</p>

Ajinkya_Ambatwar: <p>Unfortunately I couldn’t find one. Hence I moved to SGD optimizer and it worked well for me.</p>

xvdp: <p>hey, I know that this is a super late comment, just a quick note on the code; although not in this case, it could in the more general case lead to confusion. Storing transient checkpoints ought to be done with deepcopy or similar if the model is further trained.</p>
<pre><code class="lang-python">   ckp = {"optimizer":deepcopy(optimizer.state_dict()), "model":deepcopy(model.state_dict())}
</code></pre>

Noman-Tanveer: <p>Thank you, I was having an issue with the devices. This solved it for me.</p>

Manisha_Patravali: <p>Thanks a ton! This solved my pb.</p>