Restarting Optimizer and Scheduler with different learning rate

Initially, I started optimizer at LR=2e-4, and StepLR scheduler with decay of 0.1 every 50 epochs.

Now after 50 epochs, I think my LR is too high, and I want the optimizer to go to (1e-4 x 0.1) instead of (2e-4 x 0.1).

So if I change the param_groups['lr'] in the optimizer, and the scheduler.base_lrs correctly, I should be fine? I’m using Adam, and I know momentum variables might need to be handled, but otherwise, is this sufficient? (And if someone could guide me on how to adjust momentum here. I think my only choice is to reset momentum and allow the optimizer to recompute it…)

Hi,

I tried to reproduce your insight using a convNet and SGD opt as follows:

model = Net()
random_input = torch.randn(1,3,3,3)
random_target = torch.randn(3,3,3)

opt = torch.optim.SGD(model.parameters(), lr=2e-4)
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=1, gamma=0.1)
criterion = nn.MSELoss()
print(lr_scheduler.state_dict())

for step in range(1, 11):
    if step == 6:
        opt.param_groups[0]['lr'] = 1e-4
        lr_scheduler.base_lrs=[1e-4]
        print('change', lr_scheduler.state_dict())

    lr_scheduler.step(step)
    print('lr_scheduler: ',lr_scheduler.get_lr())
    print('optimizer_lr: 'opt.param_groups[0]['lr'])
    opt.zero_grad()
    output = model(random_input)
    loss = criterion(output, random_target)
    loss.backward()
    opt.step()

And it display:

{'base_lrs': [0.0002], 'gamma': 0.1, 'step_size': 1, 'last_epoch': -1}
lr_scheduler:  [2e-05]
optimizer_lr:  2e-05
lr_scheduler:  [2.0000000000000003e-06]
optimizer_lr:  2.0000000000000003e-06
lr_scheduler:  [2.0000000000000004e-07]
optimizer_lr:  2.0000000000000004e-07
lr_scheduler:  [2.0000000000000004e-08]
optimizer_lr:  2.0000000000000004e-08
lr_scheduler:  [2.0000000000000005e-09]
optimizer_lr:  2.0000000000000005e-09
change {'base_lrs': [0.0001], 'gamma': 0.1, 'step_size': 1, 'last_epoch': 5}
lr_scheduler:  [1.0000000000000004e-10]
optimizer_lr:  1.0000000000000004e-10
lr_scheduler:  [1.0000000000000004e-11]
optimizer_lr:  1.0000000000000004e-11
lr_scheduler:  [1.0000000000000006e-12]
optimizer_lr:  1.0000000000000006e-12
lr_scheduler:  [1.0000000000000005e-13]
optimizer_lr:  1.0000000000000005e-13
lr_scheduler:  [1.0000000000000006e-14]
optimizer_lr:  1.0000000000000006e-14

But it doesn’t work by only change opt.param_group[0]['lr'] and scheduler.base_lrs.
From the source code, last_epoch in lr_scheduler is updated equal to the step when we call lr_scheduler.step(step) and use it to get_lr() to update param_group[0]['lr'] in optimizer.
So, I think we should restart step in each epoch or modify learing_rate update strategy as this:

...
for step in range(1, 11):
    if step == 6:
        ...
    if step < 6:
        lr_scheduler.step(step)
    else:
        lr_scheduler.step(step-5)
    ...

if you have some other solutions, please let me know, thank you.

1 Like

Thanks for sharing! it works for me, but it seems a deprecation option to put a different epoch value in the step function.

UserWarning: The epoch parameter in `scheduler.step()` was not necessary and is being deprecated where possible. Please use `scheduler.step()` to step the scheduler. During the deprecation, if epoch is different from None, the closed form is used instead of the new chainable form, where available. Please open an issue if you are unable to replicate your use case: https://github.com/pytorch/pytorch/issues/new/choose