TopolearnA good schedule (often warmup plus cosine decay) gives better final accuracy than any fixed learning rate.
A learning-rate scheduler wraps an optimizer and adjusts each param group's lr when you call scheduler.step(). Call it after optimizer.step(), either once per epoch or once per batch, depending on how the scheduler was configured.
Common choices in torch.optim.lr_scheduler:
StepLR(opt, step_size, gamma): multiply bygammaeverystep_sizesteps.CosineAnnealingLR(opt, T_max): smooth decay along a cosine curve to near zero.OneCycleLR(opt, max_lr, total_steps): warm up then anneal, stepped every batch; popular for fast training.LinearLRfor warmup, combined with another schedule viaSequentialLR.ReduceLROnPlateau(opt): lower the rate when a monitored metric stops improving; callscheduler.step(val_loss).
Transformers are almost always trained with a short linear warmup (a few hundred to a few thousand steps) followed by cosine or linear decay. Warmup avoids unstable early updates while Adam's statistics are still poor.
warm = LinearLR(opt, start_factor=0.1, total_iters=500)
cos = CosineAnnealingLR(opt, T_max=9500)
sched = SequentialLR(opt, [warm, cos], milestones=[500])scheduler.get_last_lr() returns the current rates for logging, and schedulers have a state_dict() to save in checkpoints.
Example
import torch
from torch import nn
model = nn.Linear(2, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
sched = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.5)
for epoch in range(6):
# ... one epoch of training, calling opt.step() per batch ...
opt.step()
print(epoch, sched.get_last_lr())
sched.step() # after the optimizer step
# 0 [0.1]
# 1 [0.1]
# 2 [0.05]
# 3 [0.05]
# 4 [0.025]
# 5 [0.025]Remember
- Call
scheduler.step()afteroptimizer.step() - Warmup plus cosine decay is the modern default
- Save the scheduler's state_dict with checkpoints
Watch out for
- Stepping a per-batch scheduler (like OneCycleLR) only once per epoch, so the schedule barely moves.