Quickly Visualize PyTorch Learning Schedulers

scheduler_visualize.py

import torch
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.optim import SGD
import matplotlib.pyplot as plt

STEPS = 100

optimizer = SGD([torch.tensor(1)], lr=1)
# Use a scheduler of your choice below.
# Great for debugging your own schedulers!
scheduler = CosineAnnealingLR(optimizer, STEPS)

lrs = []
for _ in range(100):
    optimizer.step()
    lrs.append(scheduler.get_lr())
    scheduler.step()

plt.plot(lrs)
plt.show()

Thanks for the code!