Learning Rate Schedulers and Warmup

A constant learning rate is rarely optimal

Learning rate decay — starting high, ending low — is one of the simplest free wins in deep learning. PyTorch ships several schedulers; you'll use three regularly:

StepLR — multiplies LR by gamma every step_size epochs. Simple, predictable, the historical default.
CosineAnnealingLR — smooth cosine curve from initial LR to eta_min over T_max steps. The modern Transformer / vision favorite.
OneCycleLR — warmup-then-cooldown in one cycle. Excellent for fast training (Leslie Smith's 1cycle policy).

Warmup — why and how

Modern Transformer training nearly always starts with a few hundred steps of warmup: linearly increasing LR from 0 to the target, then decaying. Without it, the early gradients can be huge (parameters are random, loss is high, derivatives are wild) and the optimizer takes catastrophic steps. With warmup, you ease the optimizer into the regime where the loss landscape is well-behaved.

For warmup + cosine, you can either compose two schedulers (SequentialLR) or write a custom LambdaLR. Hugging Face's get_cosine_schedule_with_warmup is the canonical implementation if you want to skip writing it.

Step the scheduler at the right granularity

Per-epoch schedulers (StepLR, CosineAnnealingLR with epoch units) — call scheduler.step() once per epoch.
Per-step schedulers (OneCycleLR, LambdaLR sized in steps) — call scheduler.step() once per batch.

Code

StepLR and CosineAnnealingLR·python

import torch
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR

opt = optim.AdamW([torch.zeros(1, requires_grad=True)], lr=1e-3)

# StepLR: 1e-3 → (epoch 30) → 1e-4 → (epoch 60) → 1e-5
step_sched = StepLR(opt, step_size=30, gamma=0.1)

# CosineAnnealingLR: smooth decay from 1e-3 to 1e-6 over 100 epochs
cos_sched = CosineAnnealingLR(opt, T_max=100, eta_min=1e-6)

# Per-EPOCH schedulers — call once per epoch
for epoch in range(num_epochs):
    train_one_epoch(...)
    cos_sched.step()

OneCycleLR — fast training favorite·python

import torch
import torch.optim as optim
from torch.optim.lr_scheduler import OneCycleLR

steps_per_epoch = len(train_loader)
opt = optim.AdamW(model.parameters(), lr=1e-4)

scheduler = OneCycleLR(
    opt,
    max_lr=1e-3,                       # peak LR
    steps_per_epoch=steps_per_epoch,
    epochs=10,
    pct_start=0.1,                     # warmup is the first 10% of training
    anneal_strategy='cos',
)

# Per-STEP scheduler — call inside the batch loop
for epoch in range(10):
    for x, y in train_loader:
        opt.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        opt.step()
        scheduler.step()                # per batch

Warmup + cosine — the Transformer recipe·python

import math
import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR

opt = optim.AdamW(model.parameters(), lr=1e-4)
warmup_steps = 1000
total_steps = 50_000

def lr_lambda(step):
    if step < warmup_steps:
        return step / warmup_steps
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    return 0.5 * (1.0 + math.cos(math.pi * progress))

scheduler = LambdaLR(opt, lr_lambda)

# Per-step
for epoch in range(num_epochs):
    for batch in train_loader:
        opt.zero_grad()
        loss = criterion(model(batch), target)
        loss.backward()
        opt.step()
        scheduler.step()

Learning Rate Schedulers and Warmup

A constant learning rate is rarely optimal

Warmup — why and how

Step the scheduler at the right granularity

Code

External links

Exercise

Progress

Comments 0