Save the scheduler before or after the step?

1

Hello everybody,
I’m creating a script in order to save a model and, then, load the checkpoint in a second moment to resume the training.

It works very well, but I have a doubt on the point in which I have to save the scheduler. I write here the two codes and could you tell me which one is correct?

  1. Save the scheduler before the step function
for e in range(start_epoch, epochs):
    train(...)
    val(...)
    save_checkpoint(...)
    scheduler.step()
  1. Save the scheduler after the step function
for e in range(start_epoch, epochs):
    train(...)
    val(...)
    scheduler.step()
    save_checkpoint(...)

With the save_checkpoint() function, I save the model, the number of the epoch, the optimizer and the scheduler.

Thank you very much for your help.

2

I think it might depend on how you are planning to restore the model and continue the training.
If you restore the model, optimizer, and scheduler and rerun the same code, you would skip the scheduler.step() in the first example.
If you look at the complete training, you would thus use the same learning rate in two epochs, which might not be the desired use case.