Learning Rate Scheduler — Module 6.2

How This Lesson Fits the Module

A fixed learning rate rarely works for the entire training run: large rates help early exploration; small rates fine-tune near minima. Learning rate schedulers adjust LR automatically by epoch or validation feedback.

Module 6.1 introduced the learning rate hyperparameter; this lesson automates its schedule inside the training loop.

Module 6.1 Bridge See Learning Rate for the base concept. Schedulers call into the optimizer each epoch (or step) to update param_groups[0]["lr"].

Learning Objectives

By the end of this lesson, students should be able to:

Attach a PyTorch scheduler to an optimizer and call scheduler.step() correctly.
Use StepLR, ReduceLROnPlateau, and CosineAnnealingLR.
Distinguish schedulers stepped per epoch vs per batch.
Log current LR each epoch for debugging plateau behavior.
Save and restore scheduler_state_dict in checkpoints.
Match scheduler choice to training duration and architecture scale.

Common PyTorch Schedulers

Scheduler	When LR changes	Best for
`StepLR`	Every N epochs, multiply by gamma	Simple staged decay
`ReduceLROnPlateau`	When val metric stalls	Adaptive fine-tuning
`CosineAnnealingLR`	Smooth cosine decay per epoch	Modern CNN/ViT training
`OneCycleLR`	Per batch (warmup + decay)	Fast convergence (see Warmup lesson)

StepLR and CosineAnnealingLR

Time-based schedulers step once per epoch after training (and usually after validation). Read LR from the optimizer for logging.

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Option A: drop LR by 10x every 30 epochs
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

# Option B: cosine decay to eta_min over T_max epochs
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=num_epochs, eta_min=1e-6
)

for epoch in range(1, num_epochs + 1):
    train_one_epoch(...)
    val_loss, _ = validate(...)
    scheduler.step()  # epoch-based schedulers
    current_lr = optimizer.param_groups[0]["lr"]
    print(f"Epoch {epoch} | lr={current_lr:.2e} | val_loss={val_loss:.4f}")

ReduceLROnPlateau

Metric-driven scheduler steps on validation loss, not blindly by epoch. Pass the metric to scheduler.step(val_loss).

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.5,
    patience=5,
    min_lr=1e-7,
    verbose=True,
)

for epoch in range(1, num_epochs + 1):
    train_one_epoch(...)
    val_loss, _ = validate(...)
    scheduler.step(val_loss)  # pass metric, NOT epoch-based step()

Critical Mistake — Wrong scheduler.step() Signature

ReduceLROnPlateau.step(val_loss) needs the metric. Calling it with no args does nothing useful. Conversely, do not pass val_loss to StepLR—it takes no arguments.

Checkpointing Scheduler State

Resumed training must restore scheduler internal counters (last epoch, bad epochs in plateau). Include scheduler.state_dict() in checkpoints.

Engineering Habit — Log LR Every Epoch

When validation stalls, the first question is “what was the learning rate?” TensorBoard/WandB LR curves catch silent scheduler bugs.

Knowledge Check

Short Answer: Why decay learning rate during training? Answer: Large steps early, small steps for fine convergence near minima.
True/False: ReduceLROnPlateau steps every epoch regardless of metrics. Answer: False—steps when metric stops improving.
Multiple Choice: CosineAnnealingLR LR shape: (a) linear, (b) cosine, (c) random. Answer: (b).
Short Answer: Where do you read current LR? Answer: optimizer.param_groups[0]["lr"].
Short Answer: StepLR(step_size=10, gamma=0.1) effect? Answer: Multiply LR by 0.1 every 10 epochs.
True/False: Scheduler state should be in checkpoints. Answer: True when resuming training.
Multiple Choice: After validation, for StepLR you call: (a) scheduler.step(), (b) scheduler.step(val_loss), (c) optimizer.zero_grad(). Answer: (a).
Short Answer: What does factor=0.5 mean in ReduceLROnPlateau? Answer: Halve LR when triggered.
Short Answer: eta_min in cosine annealing? Answer: Minimum LR at end of cycle.
Multiple Choice: OneCycleLR steps per: (a) epoch only, (b) batch, (c) never. Answer: (b).

Key Takeaways

Schedulers modulate LR by time (Step, Cosine) or validation feedback (Plateau).
Call the correct step() signature; log LR every epoch.
Persist scheduler state when checkpointing long runs.
Next: Warmup—ramp LR up before decay.

Trainer’s Guide

Hands-on idea: Train the same model with constant LR vs cosine schedule; overlay validation curves.

Discussion prompt: When is a fixed LR acceptable (e.g. fine-tuning last layer only)?

What’s Next Many schedules start with a Warmup phase before decay.