PyTorch Learning Rate Scheduling

A learning rate scheduler adjusts the learning rate automatically during training. Starting with a larger learning rate helps the model make fast progress early. Reducing it later helps the model settle into a precise minimum instead of bouncing around it. Most state-of-the-art training recipes use a scheduler.

Why the Learning Rate Should Change

Fixed learning rate:
  Early training: lr=0.001 → good, makes fast progress
  Late training:  lr=0.001 → too large, overshoots minimum, loss oscillates

With scheduler:
  Early:  lr=0.001 → fast progress
  Middle: lr=0.0003 → refinement
  Late:   lr=0.00005 → precise convergence

Result: better final accuracy, more stable training

Step 1: Create Optimizer and Scheduler

import torch
import torch.nn as nn
import torch.optim as optim

model     = nn.Linear(10, 1)
optimizer = optim.Adam(model.parameters(), lr=0.01)

StepLR — Reduce Every N Epochs

Multiplies the learning rate by gamma every step_size epochs.

from torch.optim.lr_scheduler import StepLR

scheduler = StepLR(optimizer, step_size=10, gamma=0.5)
# lr × 0.5 every 10 epochs: 0.01 → 0.005 → 0.0025 → ...

for epoch in range(40):
    # ... training ...
    scheduler.step()     # must call once per epoch
    print(scheduler.get_last_lr())

ReduceLROnPlateau — React to Validation Loss

Reduces the learning rate when a monitored metric stops improving. This is the most practical scheduler for most projects — it adapts automatically without requiring you to guess when to reduce the rate.

from torch.optim.lr_scheduler import ReduceLROnPlateau

scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',        # watching for val_loss to go down
    factor=0.5,        # multiply lr by 0.5 when triggered
    patience=5,        # wait 5 epochs before reducing
    verbose=True
)

for epoch in range(100):
    # ... training ...
    val_loss = evaluate(model, val_loader)
    scheduler.step(val_loss)    # pass the monitored metric

Behavior:
  Epochs 1–5:  val_loss improves  → lr unchanged (0.01)
  Epochs 6–10: val_loss flat      → patience countdown (5/5)
  Epoch 11:    patience exceeded  → lr × 0.5 = 0.005
  Epochs 12–16: flat again        → countdown again
  Epoch 17:    reduce again       → lr = 0.0025

CosineAnnealingLR — Smooth Cosine Decay

The learning rate follows a cosine curve from the initial value down to a minimum. This smooth decay avoids sharp drops and works well in combination with warm restarts. It is widely used in image classification and NLP tasks.

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)
# Decays from 0.01 → 1e-6 over 50 epochs following cosine curve

for epoch in range(50):
    # ... training ...
    scheduler.step()

Cosine decay curve:
  lr
  0.01 ─┐
        │  ╲
        │    ╲
        │      ╲
        │        ╲
  1e-6  └──────────→ epoch
        0          50

OneCycleLR — Warmup + Decay

OneCycleLR ramps the learning rate up from a small value to a maximum, then decays it to a very small value. The warmup phase stabilizes early training. The decay phase refines the model. This is the default scheduler used in fast.ai and many modern training recipes.

from torch.optim.lr_scheduler import OneCycleLR

scheduler = OneCycleLR(
    optimizer,
    max_lr=0.01,
    steps_per_epoch=len(train_loader),
    epochs=30
)

for epoch in range(30):
    for batch_X, batch_y in train_loader:
        # ... training step ...
        scheduler.step()    # called per BATCH, not per epoch

OneCycle lr curve:
  lr
  0.01   ─────╮
              │ ╲
              │   ╲
  0.001 ─╯    │     ╲
  Warmup  Peak  Annealing

Logging the Learning Rate

current_lr = optimizer.param_groups[0]['lr']
print(f"Epoch {epoch}: lr = {current_lr:.6f}")

Scheduler Placement

Call scheduler.step():
  → After each EPOCH for: StepLR, ReduceLROnPlateau, CosineAnnealingLR
  → After each BATCH for: OneCycleLR

Call order:
  optimizer.step()     ← update weights FIRST
  scheduler.step()     ← then update lr

Scheduler Comparison

Scheduler           Best For
───────────────────────────────────────────────────────
StepLR              Simple, predictable decay
ReduceLROnPlateau   Practical, metric-driven — good default
CosineAnnealingLR   Smooth decay, image tasks
OneCycleLR          Fast training, state-of-the-art recipes

Summary

Learning rate schedulers automatically adjust the learning rate during training for better convergence. StepLR reduces by a fixed factor every N epochs. ReduceLROnPlateau reacts to validation loss stopping its improvement — the most practical choice. CosineAnnealingLR follows a smooth cosine curve. OneCycleLR ramps up then decays and is called per batch. Call scheduler.step() after optimizer.step(). For ReduceLROnPlateau, pass the validation loss as an argument. Log the current learning rate each epoch to verify the scheduler is behaving as expected.