Using steps rather than epochs for learning rate scheduler

[cmlakhan]

I'd like to use the transformers.get_constant_schedule_with_warmup learning rate scheduler from HuggingFace but I get the impression that when I set the "num_warmup_steps" value it corresponds to epochs rather than steps. What is the proper way to change this from epochs to steps? This is the way I have it coded

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=2e-3, weight_decay=0.000001)
        warmup_scheduler = transformers.get_constant_schedule_with_warmup(optimizer, num_warmup_steps=5000)
        optimizer_dict = {'optimizer': optimizer, 'lr_scheduler': warmup_scheduler}
        return optimizer_dict

[akihironitta]

Hi @cmlakhan! You can call lr scheduler's step() every step (by default, it's called every epoch as you're aware already) by passing a dictionary to the "lr_scheduler" field.

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=2e-3, weight_decay=0.000001)
        warmup_scheduler = transformers.get_constant_schedule_with_warmup(optimizer, num_warmup_steps=5000)
+       lr_scheduler_config = {'scheduler': warnup_scheduler, 'interval': 'step'}
+       optimizer_dict = {'optimizer': optimizer, 'lr_scheduler': lr_scheduler_config)
-       optimizer_dict = {'optimizer': optimizer, 'lr_scheduler': warmup_scheduler}
        return optimizer_dict

Also, we have a few examples in our documentation at: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#configure-optimizers