Continuing training when using learning rate schedulers

[KirillShmilovich]

❓ Questions and Help

When restarting training on a model using learning rate scheduler, it seems like the original learning rate is used rather than the scheduler-update learning rate.

Code

For example, a model with the following configure_optimizers:

def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        scheduler = optim.lr_scheduler.ExponentialLR(
            optimizer, gamma=self.hparams.learning_gamma
        )
        return [optimizer], [scheduler]

With learning_gamma != 1.0, when restarting training, e.g.:

model = myModel.load_from_checkpoint(ckpt_fname)
lr_monitor = LearningRateMonitor(logging_interval="epoch")
trainer = Trainer(resume_from_checkpoint=ckpt_fname, callbacks=[lr_monitor])
trainer.fit(model)

The logged learning rate is equal to the original initial learning rate, rather than the schedule-updated learning rate.

What's your environment?

• OS: Linux
• Packaging: conda
• Version: 1.1.0

[rohitgr7]

yeah, it's a bug. Fixing it here: #5208.. probably will complete it this week.

you mentioned that it starts from the beginning. Is your resume_from_checkpoint contain the state trained for a single epoch?

[KirillShmilovich]

Thank you for your help with this!

Here is the callback I use to generate my checkpoint during training:

class CheckpointEveryNSteps(pl.Callback):
    """
    Save a checkpoint every N steps, instead of Lightning's default that checkpoints
    based on validation loss.
    """

    def __init__(
        self,
        save_step_frequency,
        prefix="N-Step-Checkpoint",
        use_modelcheckpoint_filename=False,
    ):
        """
        Args:
            save_step_frequency: how often to save in steps
            prefix: add a prefix to the name, only used if
                use_modelcheckpoint_filename=False
            use_modelcheckpoint_filename: just use the ModelCheckpoint callback's
                default filename, don't use ours.
        """
        self.save_step_frequency = save_step_frequency
        self.prefix = prefix
        self.use_modelcheckpoint_filename = use_modelcheckpoint_filename

    def on_batch_end(self, trainer: pl.Trainer, _):
        """ Check if we should save a checkpoint after every train batch """
        # epoch = trainer.current_epoch
        global_step = trainer.global_step
        if global_step % self.save_step_frequency == 0:
            if self.use_modelcheckpoint_filename:
                filename = trainer.checkpoint_callback.filename
            else:
                # filename = f"{self.prefix}_epoch={epoch}_global_step={global_step}.ckpt"
                filename = f"{self.prefix}.ckpt"
            ckpt_path = os.path.join(trainer.checkpoint_callback.dirpath, filename)
            trainer.save_checkpoint(ckpt_path)

[rohitgr7]

hmm... the existing issue I am aware of is with epoch-level schedulers. Mind reproduce it with https://colab.research.google.com/drive/1HvWVVTK8j2Nj52qU4Q4YCyzOm0_aLQF3?usp=sharing.. will check.

[KirillShmilovich]

As far as I know optim.lr_scheduler.ExponentialLR is an epoch-level scheduler? https://pytorch.org/docs/stable/optim.html#torch.optim.lr_scheduler.ExponentialLR

[rohitgr7]

yeah.... didn't check the scheduler you were using. That's why I asked, did you try it with any checkpoint that is saved after the 2nd epoch is completed? because at this point, the checkpoint lags with just one scheduler step.

[KirillShmilovich]

Yes, I've tried this after many epochs (10+). Regardless, the learning rate always resets to it's initial value.

[rohitgr7]

mind reproduce it with https://colab.research.google.com/drive/1HvWVVTK8j2Nj52qU4Q4YCyzOm0_aLQF3?usp=sharing... will check it closely what's happening here.

[KirillShmilovich]

https://colab.research.google.com/drive/1HF251ZR4we6l21BCOgfG1A6Ay9TM7ba4?usp=sharing

Looks like this actually isn't a problem... Not sure why LearningRateMonitor was logging incorrect values. I'll make sure things are working on my end too...

[rohitgr7]

great, pls reopen if it's still a problem. BTW the notebook you shared is not public.