How to use in `pytorch-lightning`? about pytorch_warmup

[Tony-Y]

I don't know how to use pytorch_warmup inside pytorch-lightning. But you could use it by reading pytorch-lightning docs: Bring your own Custom Learning Rate Schedulers

import torch
import pytorch_warmup as warmup


def configure_optimizers(self):
    optimizer = torch.optim.AdamW(self.model..parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=0.01)
    num_steps = len(self.train_dataloader) * self.num_epochs # I don't know whether this line is correct or not.
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps)  
    warmup_scheduler = warmup.UntunedLinearWarmup(optimizer)
    return [optimizer], [{"scheduler": (lr_scheduler, warmup_scheduler), "interval": "step"}]


def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
    lr_scheduler, warmup_scheduler = scheduler
    with warmup_scheduler.dampening():
            lr_scheduler.step()

[Tony-Y]

I have just done it on Colab:
PyTorch Lightning with pytorch_warmup.ipynb