Implementation of get_linear_schedule_with_warmup in lightning-hydra-template

[yipliu]

configs/model/mnist.yaml

scheduler:
  _target_: torch.optim.lr_scheduler.LambdaLR
  _partial_: true
  last_epoch: -1

src/models/mnist_module/MNISTLitModule

def configure_optimizers(self):
    def lr_lambda(current_step):
        num_cycles = 0.5
        num_training_steps = self.trainer.estimated_stepping_batches
        if current_step < self.hparams.num_warmup_steps:
            return float(current_step) / float(max(1, self.hparams.num_warmup_steps))
        progress = float(current_step - self.hparams.num_warmup_steps) / float(
                max(1, num_training_steps - self.hparams.num_warmup_steps)
            )
        return progress
    optimizer = self.hparams.optimizer(params=self.parameters())
    scheduler = self.hparams.scheduler(optimizer=optimizer, lr_lambda=lr_lambda)
return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            },
        }