How do you implement and use learning rate schedulers in PyTorch?

A fixed learning rate throughout training is rarely optimal — too high late in training prevents fine convergence, while too low early on wastes time. PyTorch's torch.optim.lr_scheduler module adjusts the learning rate systematically as training progresses.

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

model     = nn.Linear(10, 1)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)

# ── StepLR: multiply lr by gamma every step_size epochs
scheduler_step = optim.lr_scheduler.StepLR(
    optimizer, step_size=10, gamma=0.1
)

# ── CosineAnnealingLR: smooth decay following a cosine curve
scheduler_cos = optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=100, eta_min=1e-6
)

# ── ReduceLROnPlateau: reduce lr when a metric stops improving
scheduler_plateau = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.5, patience=5
)

# ── OneCycleLR: warmup then decay — fast convergence ("super-convergence")
n_epochs, steps_per_epoch = 10, 100
scheduler_1cycle = optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=1e-2,
    total_steps=n_epochs * steps_per_epoch,
    pct_start=0.3,    # 30% of steps used for warmup
)

# ── Training loop with scheduler
for epoch in range(100):
    train_one_epoch(model, train_loader, optimizer, loss_fn, device)
    val_loss, _ = validate(model, val_loader, loss_fn, device)

    scheduler_cos.step()              # epoch-based scheduler — call once per epoch
    scheduler_plateau.step(val_loss)  # metric-based — pass the metric value

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

# Note: OneCycleLR and some schedulers are called PER BATCH, not per epoch
# for step in range(total_steps):
#     train_step(...)
#     scheduler_1cycle.step()  # called inside the batch loop

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.