How to freely change LR scheduler during training

[LLNLanLeN]

I have a question regarding Trainer and LightningModule. I want to be able to freely change the LR scheduler (while keeping the Optimizer the same) during the training step. This may have been answered before but I haven't been able to find a working solution so far. I found one thread that related the most to my question, but unfortunately I could not get it to work:

#3095

My original idea was similar to what suggested in the issue below, to create a simple callbacks and swap out the old LR Scheduler with a new one (just LR scheduler, I want to keep the optimizer the same). But I didn't have much success trying the solution above.

Note: the Pytorch Lightning version that I'm using is 1.1.8, and not the latest, which is 1.4.0. Some of the custom packages that I'm using is more compatible with 1.1.8 instead of 1.4.0 so I opt for the older version instead.

This is my implementation:

class ConfigureMultiStepLRScheduler(Callback):
    def __init__(self,
                 learning_rate_decay: float,
                 learning_rate_milestone: List[int]):
        super().__init__()

        self.learning_rate_decay = learning_rate_decay
        self.learning_rate_milestone = learning_rate_milestone

    def on_train_epoch_start(self, trainer, pl_module):
        new_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer = trainer.optimizers,
                                                             milestones = self.learning_rate_milestone,
                                                             gamma = self.learning_rate_decay)
        trainer.lr_schedulers = trainer.configure_schedulers([new_scheduler])
....
....
....
....
callbacks = [
        ConfigureMultiStepLRScheduler(learning_rate_decay = LEARNING_RATE_DECAY,
                                      learning_rate_milestone = LEARNING_RATE_MILESTONE),
    ]

    trainer = pl.Trainer(
        gpus = 1,
        accelerator = "ddp",
        max_epochs = TOTAL_EPOCH,
        callbacks = callbacks,
        check_val_every_n_epoch = 1,
        limit_train_batches = 5,
        deterministic = True,
    )

I ran into this error with my implementation:

Traceback (most recent call last):
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/lanle/PHIAR/ai-pipeline/quantization/AIMET_semantic_segmentation_QAT.py", line 520, in <module>
    main(args)
  File "/home/lanle/PHIAR/ai-pipeline/quantization/AIMET_semantic_segmentation_QAT.py", line 499, in main
    dm.val_dataloader())
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 510, in fit
    results = self.accelerator_backend.train()
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 158, in train
    results = self.ddp_train(process_idx=self.task_idx, model=model)
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 307, in ddp_train
    results = self.train_or_test()
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/accelerators/accelerator.py", line 74, in train_or_test
    results = self.trainer.train()
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 557, in train
    self.train_loop.on_train_epoch_start(epoch)
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 230, in on_train_epoch_start
    self.trainer.call_hook("on_train_epoch_start")
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 926, in call_hook
    trainer_hook(*args, **kwargs)
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/pytorch_lightning/trainer/callback_hook.py", line 72, in on_train_epoch_start
    callback.on_train_epoch_start(self, self.get_model())
  File "/home/lanle/PHIAR/ai-pipeline/quantization/AIMET_semantic_segmentation_QAT.py", line 223, in on_train_epoch_start
    gamma = self.learning_rate_decay)
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/torch/optim/lr_scheduler.py", line 382, in __init__
    super(MultiStepLR, self).__init__(optimizer, last_epoch)
  File "/home/lanle/anaconda3/envs/AIMET/lib/python3.6/site-packages/torch/optim/lr_scheduler.py", line 30, in __init__
    type(optimizer).__name__))
TypeError: list is not an Optimizer