Create 2 learning rate schedulers for 1 optimizer

[davidnvq]

In my custom pl.LightningModule, I try to specify two learning rate schedulers (warming up in the first epoch, and multi-step scheduler in the following epochs) for 1 optimizer as follows:

    def configure_optimizers(self):
        args = self.args

        optimizer = torch.optim.SGD(self.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)

        warmup_iters = min(1000, self.num_batches - 1)
        warmup_scheduler = {
                    'scheduler': utils.warmup_lr_scheduler(optimizer, warmup_iters, 1./1000),
                    'name': 'learning_rate',
                    'interval':'step',
                    'frequency': 1
                }

        lr_scheduler = {
            'scheduler': torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma),
            'name': 'learning_rate',
            'interval':'epoch',
            'frequency': 1
        }
        return [optimizer], [warmup_scheduler, lr_scheduler]

However, it seems like the learning rate in the optimizer only controlled by the first warmup_scheduler.

So how can I create 2 learning rate schedulers for 1 optimizer?
Many thanks.

[davidnvq]

The answer of my question would be involved and solve the question in #6234.
Thus, I am still looking for the answer :)

[tchaton]

Dear @davidnvq,

I don't believe there is a simple out-of-box solution in Lightning right now.

Here are the 2 suggested approach:

• Implement a callback which handles lr_scheduler stepping.
• Fuse your 2 scheduler into 1.

[thesofakillers]

Could you maybe use a ChainedScheduler?