pytorch-lightning Transformer (noam) lr policy

model_with_noam.py

import torch
import pytorch_lightning as pl

class MyTransformer(pl.LightningModule):
    def __init__(
        self,
        learning_rate=0.001,
        warmup=4000,
    ):
      self.learning_rate = learning_rate
      self.warmup = warmup

    #
    # This is an example config on how to do the Noam (Attention is all you need) lr policy
    # with pytorch lightning.
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)

        def warm_decay(step):
            if step < self.warmup:
                return  step / self.warmup
            return self.warmup ** 0.5 * step ** -0.5

        scheduler = (
            {
                "scheduler": torch.optim.lr_scheduler.LambdaLR(optimizer, warm_decay),
                "interval": "step", #runs per batch rather than per epoch
                "frequency": 1,
                # "name" : "learning_rate" # uncomment if using LearningRateMonitor
            }
        )

        return [optimizer], [scheduler]