Implementing a Custom Learning Rate Scheduler

medium (<1 hr) training custom scheduler learning rate

this month by E

Implement a custom learning rate scheduler that follows a cosine annealing schedule. The learning rate starts high and decreases smoothly to a minimum value, then resets and repeats. Your scheduler should be a subclass of torch.optim.lr_scheduler.LRScheduler and should update the learning rate based on the current epoch.

$l r_{t} = l r_{m i n} + 0.5 * (l r_{m a x} - l r_{m i n}) * (1 + c o s (t / T * π))$ where $t$ is the current epoch and $T$ is the total number of epochs for one cycle.

Verification: Train a small model for several epochs and print the learning rate at each epoch. The values should follow the expected cosine curve, with resets at the end of each cycle.