Restarting with adam

Question

I am training my network with early stopping strategy. I start with a higher learning rate, and based on validation loss, I need to restart training from an earlier snapshot.

I am able to save/load snapshot with model and optimizer state_dicts . No problem with that.

My question is, once I restart training, how do I set the learning rate of adam again? Should I restart adam fresh instead of using a state_dict or should I use optimizer.param_groups[0]['lr'] = lr to adjust learning rate with loaded optimizer state_dict ?

For example, I train my network with lr = 1e-6 for 5 epochs, saved model and optimizer state_dict . I am now restarting from epoch 6, but I need lr = 1e-7 instead. What is the best approach for this?

Thanks!

Answer 1

Looking at PyTorch's torch.optim.lr_scheduler code here , I can see that they set the parameter of the optimizer. Thus, that will be the best approach. The exact place I can see this is in step function of class _LRScheduler (in the above link).

You can do the same by

optimizer.param_groups[0]['lr'] = lr

as you had mentioned yourself.

Answer 2

Looking further into the scheduler code, I found the correct way to do it as:

def get_lr(gamma, optimizer):
    return [group['lr'] * gamma
            for group in optimizer.param_groups]

for param_group, lr in zip(optimizer.param_groups, get_lr(gamma, optimizer)):
    param_group['lr'] = lr