Discontinuity in learning rate value when resuming training from checkpoint

clems: <p>hey, I’m trying to resume training from a given checkpoint using pytorch CosineAnnealingLR scheduler.<br>
let’s say I want to train a model for 100 epochs, but, for some reason, I had to stop training after epoch 45 but saved both the optimizer state and the scheduler state.<br>
I want to resume training from epoch 46. I’ve followed what has previously been chatted on this forum to resume training from a given epoch, but when plotting learning rates values as a function of epochs, I get a discontinuity at epoch 46 (see figure below, plot on the left).</p>
<p>For comparison, I run the full 100 epochs and plotted the learning rate to show what the expected plot should look like (see figure below, plot in the center).</p>
<p>We can see both plots do not match when displaying them on the same figure (see figure below, plot on the right ; in green: expected plot ; in blue: plot with discontinuity)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/8/28ff2e58269cb43ef5bc6444f15d81c8261f21a9.png" data-download-href="https://discuss.pytorch.org/uploads/default/28ff2e58269cb43ef5bc6444f15d81c8261f21a9" title="figure"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/8/28ff2e58269cb43ef5bc6444f15d81c8261f21a9_2_689x168.png" alt="figure" data-base62-sha1="5QFOWsiqWr3YJQOyd5x3p59FFj3" width="689" height="168" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/8/28ff2e58269cb43ef5bc6444f15d81c8261f21a9_2_689x168.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/8/28ff2e58269cb43ef5bc6444f15d81c8261f21a9_2_1033x252.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/8/28ff2e58269cb43ef5bc6444f15d81c8261f21a9.png 2x" data-dominant-color="F7F8F7"><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">figure</span><span class="informations">1240×302 22.4 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>Here is a snippet of the code I’ve used to resume training:</p>
<pre><code class="lang-auto">intial_epoch = 0
nepochs_first = 45
nepochs_total = 100

base_lr = 0.0001
optimizer_first = torch.optim.Adam(model.parameters(), lr=base_lr)

scheduler_first = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_first, T_max=nepochs_total, last_epoch=intial_epoch-1)

lr_first = []
for i in range(intial_epoch+1, nepochs_first+1):
    scheduler_first.step()
    lr_first.append(scheduler_first.get_last_lr()[-1])

optimizer_state, scheduler_state = optimizer_first.state_dict(), scheduler_first.state_dict()

optimizer = torch.optim.Adam(model.parameters(), lr=1)
# I deliberately set the initial lr to a different value than base_lr, and it should be overwritten when loading the state_dict
optimizer.load_state_dict(optimizer_state)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=nepochs_total, last_epoch=nepochs_first-1)
scheduler.load_state_dict(scheduler_state)

prev_epoch = scheduler_state['last_epoch']
lr = []
for i in range(prev_epoch+1, nepochs_total+1):
    scheduler.step()
    lr.append(scheduler.get_last_lr()[-1])
</code></pre>
<p>I’ve tried a bunch of things but I couldn’t manage to get over this discontinuity.<br>
Thanks you for your help!</p>

clems: <p>the discontinuity impacts all LR values computed from epoch 46 to epoch 100.<br>
indeed, see what this quick snippet gave:</p>
<pre><code class="lang-auto">for i, val in enumerate(lr_first+lr):
    if val != lr_expected[i]:
        print(f'epoch: {i+1} \t actual lr: {val} \t expected lr: {lr_expected[i]}')
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/0/0084239ca8c2fe48431253e16231395e79d95c82.png" data-download-href="https://discuss.pytorch.org/uploads/default/0084239ca8c2fe48431253e16231395e79d95c82" title="Screen Shot 2020-08-17 at 3.29.50 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/0084239ca8c2fe48431253e16231395e79d95c82_2_690x294.png" alt="Screen Shot 2020-08-17 at 3.29.50 PM" data-base62-sha1="4z6Ci7dtDYqTefW0cgcCGxaywW" width="690" height="294" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/0084239ca8c2fe48431253e16231395e79d95c82_2_690x294.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/0084239ca8c2fe48431253e16231395e79d95c82_2_1035x441.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/0084239ca8c2fe48431253e16231395e79d95c82_2_1380x588.png 2x" data-dominant-color="E2E2E2"><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">Screen Shot 2020-08-17 at 3.29.50 PM</span><span class="informations">1602×684 124 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>

clems: <p>any ideas what I did wrong?</p>

ebelle: <p>Have you tried saving the learning rate from the previous model and putting that in when you set up the optimizer instead of relying on the state dict?</p>
<p>optimizer = torch.optim.Adam(model.parameters(), lr=previously_saved_lr)</p>

clems: <p>hi, thanks for the suggestion but the learning rate scheduler returns an error if last_epoch != -1 and optimizer’s state doesn’t have a ‘initial_lr’ key</p>
<p>The issue I’m facing may come from pytorch _LRScheduler class <code>__init__</code> method, which ends with a <code>self.step()</code> call that may change the learning rate value…<br>
I’m thinking about this because as soon as I instantiate the LR scheduler, the optimizer’s learning rate is modified:</p>
<pre><code class="lang-auto"># let's say we have stopped training after epoch 45 out of 100 
# and we want to resume training
epoch_restart = 45

new_optimizer = torch.optim.Adam(model.parameters(), lr=1)
# load optimizer state saved when training stopped 
new_optimizer.load_state_dict(optimizer_state)

# retrieve optimizer's lr value before instantiating LR scheduler
lr_first = new_optimizer.state_dict()['param_groups'][0]['lr']

new_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    new_optimizer,
    T_max=100,
    eta_min=0,
    last_epoch=epoch_restart,
)

# retrieve optimizer's lr value after instantiating LR scheduler
lr_then = new_optimizer.state_dict()['param_groups'][0]['lr']
</code></pre>
<p>if you compare <code>lr_first</code> and <code>lr_then</code>, they’ll be different!<br>
why?</p>