LR scheduler steps after saving checkpoint with iteration-based checkpointing

[simran2905]

🐛 Bug

When checkpoints are saved using every_n_train_steps argument of ModelCheckpoint callback, checkpointing happens before lr_scheduler step being called. While it is the other way around with every_n_val_epochs argument.

In training_loop, on_train_batch_end (which saves checkpoint) is called before updating lr_scheduler. https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/training_loop.py#L500-L529

while on_evaluation_end is called after updating lrs. https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/trainer.py#L996-L1011

Please reproduce using the BoringModel

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

To Reproduce

Expected behavior

In the colab notebook, a training epoch is of 64 batches. When I save checkpoint using every_n_train_steps=64 or with every_n_val_epochs=1, the lr_scheduler step count should be the same. But the step count is 1 lesser with every_n_train_steps=64

Environment

• PyTorch Version (e.g., 1.0):
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

Additional context

cc @ananthsub

[tchaton]

Dear @simran2905,

Yes, that's definitely a bug. Thanks for reporting it.
Would you mind making a contribution and try to resolve this behaviour ?

Best,
T.C

[simran2905]

Would you mind making a contribution and try to resolve this behaviour ?

Hey @tchaton, I am not super familiar with the training loop in PL. So this would be difficult. Sorry.

[Lucklyric]

@simran2905 , it seems the problem has been fixed by a recent refactor on training_loop in #7357 .

[simran2905]

Hey @Lucklyric , yes the training loop refactor fixes step and epoch based checkpoints to save the same state. But I believe the lr_scheduler step should happen before checkpointing.

Here is a colab notebook demonstrating the issue. https://colab.research.google.com/drive/1x90xCxF658bve7QTLx0CEXVMWRdPTajG?usp=sharing

The lr is divided by 10 at every step. When I run it for 2 epochs (saving checkpoint at the end of epoch), lr at step 64 is 1e-67. When I resume from the checkpoint saved at epoch 1, lr at step 64 is 1e-66 because the lr update was not saved with the checkpoint.

cc @carmocca as the author of #7357

Let me know if I should file another issue for this.

[ananthsub]

@carmocca do you see any issue with moving https://github.com/PyTorchLightning/pytorch-lightning/blob/36770b22fd6370425e377a7fe94d74fd29af821b/pytorch_lightning/trainer/training_loop.py#L526-L527 to before self.on_train_batch_end ?

[simran2905]

If on_train_batch_end is logging some metrics that are later being used for Lr scheduler monitor, then this move won't work IMO. Only the checkpoint one should be called later.

[carmocca]

When I resume from the checkpoint saved at epoch 1, lr at step 64 is 1e-66 because the lr update was not saved with the checkpoint.

Thanks for the colab link!

The test for your case would be the following:

def test_scheduler_lr_step_interval_updated_before_saving(tmpdir):
    batches = 2
    lr, gamma = 1, 10
    trainer = Trainer(
        default_root_dir=tmpdir,
        progress_bar_refresh_rate=0,
        logger=False,
        max_epochs=1,
        limit_train_batches=batches,
        limit_val_batches=1,
        callbacks=[ModelCheckpoint(dirpath=tmpdir)]
    )

    class TestModel(BoringModel):

        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.parameters(), lr=lr)
            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma)
            lr_dict = {'scheduler': lr_scheduler, 'interval': 'step'}
            return [optimizer], [lr_dict]

        def on_save_checkpoint(self, checkpoint):
            lr_dict = checkpoint['lr_schedulers'][0]
            # 2 batches ran. since the lr_dict interval is `step`, the step count should be 2
            assert self.trainer.global_step + 1 == batches  # the global step hasn't been increased yet
            assert lr_dict['_step_count'] - 1 == batches  # step count starts at 1
            assert lr_dict['_last_lr'] == [lr * gamma ** batches]
            self.on_save_checkpoint_called = True

    model = TestModel()
    trainer.fit(model)
    assert model.on_save_checkpoint_called

do you see any issue with moving to before self.on_train_batch_end ?

Yes. Namely, that on_train_batch_end is allowed to self.log

https://github.com/PyTorchLightning/pytorch-lightning/blob/36770b22fd6370425e377a7fe94d74fd29af821b/pytorch_lightning/trainer/connectors/logger_connector/fx_validator.py#L55

Let me elaborate:

The previous test fails on master and passes if we swap

if should_check_val:
    self.trainer._run_evaluation()

...

self.update_lr_schedulers('step')

for

self.update_lr_schedulers('step')

if should_check_val:
    self.trainer._run_evaluation()

Because in this case, the checkpoint gets saved when validation is over so moving the scheduler update before means the checkpoint state is the expected.

This seems easy until you remember about our friend ReduceLROnPlateau. This scheduler can track any logged metric so if that metric happens to get logged during validation, the update_lr_scheduler() call will fail with ReduceLROnPlateau conditioned on metric foo which is not available. Available metrics are: []. Condition can be set using monitor key in lr scheduler dict.

This is really inconvenient since there is no way to know when an OnPlateau metric will be logged.
So we will probably have to call twice the update. Once before for the classic schedulers and once after for the OnPlateau schedulers.

Any thoughts?