SWA steps with lr scheduler step interval

[mibaumgartner]

🐛 Bug

I'm currently working on a custom implementation of SWA with cyclic learning rates and was wondering about the replacement procedure in the current implementation.

https://github.com/PyTorchLightning/pytorch-lightning/blob/2f6ce1ae7fff34d16d3707571f6a9a7b0fb0c50a/pytorch_lightning/callbacks/stochastic_weight_avg.py#L191-L195

If an lr scheduler is already configured, only the scheduler key is exchanged for the SWALR scheduler which expects to be stepped at the end of the epoch. I was wondering if this is correct if the original lr scheduler was configured with interval=step ?

Please reproduce using the BoringModel

To Reproduce

Use following BoringModel and post here

Expected behavior

The SWALR scheduler is only stepped at the end of the epoch and the configuration is exchanged too.

_scheduler_config = _get_default_scheduler_config()
assert _scheduler_config["interval"] == "epoch"
_scheduler_config["scheduler"] = self._swa_scheduler

if trainer.lr_schedulers:
    lr_scheduler = trainer.lr_schedulers[0]["scheduler"]
    rank_zero_warn(f"Swapping lr_scheduler {lr_scheduler} for {self._swa_scheduler}")
    trainer.lr_schedulers[0] = _scheduler_config
else:
    trainer.lr_schedulers.append(_scheduler_config)

Environment

Note: Bugs with code are solved faster ! Colab Notebook should be made public !

• IDE: Please, use our python bug_report_model.py template.

• Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).

You can get the script and run it with:

wget https://github.com/raw/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py

• 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

@tchaton

[tchaton]

Dear @mibaumgartner,

You are right, there is definitely a bug there. We should step SWA scheduler only on epoch.

Would you mind making a PR to rectify this behaviour ?

Best,
T.C

[mibaumgartner]

Hi @tchaton,

I'll submit a PR today :)

Best,
Michael