PyTorch learning rate scheduler: examples and step timing

A PyTorch learning-rate scheduler changes the learning rate held by an optimizer. Before choosing one, decide what advances the schedule: a completed epoch, an optimizer update, or a validation result. Placing the same call in a different part of the loop changes the training schedule.

This guide compares four common choices and provides three small runnable examples. Fourteen correctness checks passed on CPU with PyTorch 2.15.0 and Python 3.12.14, including a checkpoint continuation. These tests establish the behavior shown here, not improvements in accuracy, convergence speed, or GPU performance.

Choose the schedule and its unit

Scheduler	What controls it	Decision to make first
StepLR	A fixed interval between reductions	How often you call step; this guide uses once per epoch
CosineAnnealingLR	A cosine curve over a planned horizon	The number of calls represented by T_max
ReduceLROnPlateau	A metric supplied to step	The metric, its direction, and validation frequency
OneCycleLR	A fixed budget of training updates	The number of optimizer updates in the run

For an epoch-based schedule, call the scheduler after that epoch’s optimizer updates. For an update-based schedule, call it after the corresponding update. PyTorch’s optimizer documentation explains the required order: putting the scheduler before the optimizer can skip the first scheduled rate.

Keep the unit in the experiment configuration and logs. A scheduler does not count your data automatically. Changing batch size, dropping a final batch, or changing gradient accumulation can alter the number of updates. Our DataLoader guide covers the batching decisions behind those counts.

StepLR: see the rate used and the rate prepared next

StepLR multiplies the learning rate by gamma at each configured interval. Save this self-contained example as step_lr.py and run it in an environment with PyTorch installed:

import torch
from torch import nn

x = torch.tensor([[0.0], [1.0], [2.0], [3.0]])
y = 2 * x + 1


def build():
    torch.manual_seed(7)
    model = nn.Linear(1, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=2, gamma=0.5
    )
    return model, optimizer, scheduler


def train_epoch(model, optimizer):
    model.train()
    for start in range(0, len(x), 2):
        optimizer.zero_grad(set_to_none=True)
        loss = nn.functional.mse_loss(model(x[start:start + 2]), y[start:start + 2])
        loss.backward()
        optimizer.step()


if __name__ == "__main__":
    model, optimizer, scheduler = build()
    for epoch in range(6):
        used_lr = optimizer.param_groups[0]["lr"]
        train_epoch(model, optimizer)
        scheduler.step()
        next_lr = optimizer.param_groups[0]["lr"]
        print(f"epoch={epoch + 1} used={used_lr:.4f} next={next_lr:.4f}")

Each epoch contains two optimizer updates. The scheduler advances once after both, so step_size=2 means two epochs here. The tiny regression problem provides gradients; its settings are illustrative, not a recommendation for a real dataset.

The run produced these rates:

Epoch	Used during the epoch	Prepared for the next epoch
1	0.1000	0.1000
2	0.1000	0.0500
3	0.0500	0.0500
4	0.0500	0.0250
5	0.0250	0.0250
6	0.0250	0.0125

The final 0.0125 value was prepared but never used by this six-epoch run. Labeling those two moments prevents a common logging mistake. For several parameter groups, record every group’s lr, rather than only group zero. The scheduler API also provides get_last_lr() for its most recently calculated rates.

ReduceLROnPlateau: step after validation

ReduceLROnPlateau takes a metric: use mode="min" when lower is better, or mode="max" when higher is better. Feed it a consistently calculated validation result at a defined interval.

This separate script uses synthetic metrics to make the patience behavior visible. It does not train a model or measure validation quality:

import torch

parameter = torch.nn.Parameter(torch.tensor([0.0]))
optimizer = torch.optim.SGD([parameter], lr=0.1)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.5, patience=1,
    threshold=0.0, threshold_mode="abs"
)

# Synthetic metrics to inspect scheduler behavior; no model training here.
for validation_loss in [1.0, 0.8, 0.8, 0.8, 0.7]:
    scheduler.step(validation_loss)
    print(f"metric={validation_loss:.1f} next={optimizer.param_groups[0]['lr']:.3f}")

The rates after these calls were 0.100, 0.100, 0.100, 0.050, 0.050. With an absolute threshold of zero, the repeated 0.8 values do not improve on the best value. Patience 1 permits one such event; the second triggers the reduction. The final 0.7 improves the metric and resets the bad-event count.

In a training loop, finish validation before calling scheduler.step(validation_loss). Use the intended aggregate metric, not whichever training batch happened to run last. If validation frequency changes, the same patience value covers a different amount of training. Set threshold, cooldown, and minimum rate deliberately for the metric’s scale and noise.

OneCycleLR: count actual optimizer updates

OneCycleLR increases and then decreases the rate over a specified budget. It can take total_steps directly or infer the budget from epochs and steps per epoch. In this standalone example, one loop iteration performs one optimizer update:

import torch

parameter = torch.nn.Parameter(torch.tensor([1.0]))
optimizer = torch.optim.SGD([parameter], lr=0.1)
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=0.1, total_steps=8,
    pct_start=0.5, cycle_momentum=False
)

for update in range(8):
    optimizer.zero_grad(set_to_none=True)
    loss = parameter.square().sum()
    loss.backward()
    used_lr = optimizer.param_groups[0]["lr"]
    optimizer.step()
    scheduler.step()
    print(f"update={update + 1} used={used_lr:.8f}")

The measured rates used by the first, fourth, and eighth updates were 0.004, 0.1, and 0.0000004. Constructing the scheduler sets the initial rate, so the optimizer’s earlier lr=0.1 is not the first rate used. cycle_momentum=False keeps this demonstration focused on learning rates. A ninth scheduler call raised an error because the eight-step budget was exhausted.

With gradient accumulation, several microbatches may produce one optimizer update. Count that update once when using an update-based schedule, and calculate the budget from your actual loop. Our gradient-accumulation guide covers the update boundary and partial final groups.

Mixed precision adds another case: PyTorch’s AMP examples explain that a scaler can skip an optimizer step when gradients contain nonfinite values. Keep an update-based scheduler aligned with successful updates in that integration. The CPU tests here do not validate AMP or distributed training.

CosineAnnealingLR: define the horizon explicitly

CosineAnnealingLR uses T_max to define its horizon and eta_min as the minimum rate. Its counter advances with scheduler calls. If you call it after every batch instead of every epoch, the same T_max describes a different amount of training.

A separate CPU check starting at 0.1 with T_max=4 and eta_min=0.01 reached 0.01 after four calls. This scheduler does not implement warm restarts. Plan the intended endpoint rather than assuming additional calls will hold the minimum indefinitely.

Resume the schedule with the training state

Save the model, optimizer, scheduler state, and progress counter at a clearly defined boundary. Reconstruct the matching objects before restoring them. PyTorch’s optimizer loading reference warns that creating the scheduler after loading optimizer state can overwrite restored learning rates.

For our StepLR check, we saved after epoch three, including the scheduler call. We reconstructed all three objects, restored model and scheduler state, then restored the optimizer. The remaining rates were 0.05, 0.025, 0.025, matching uninterrupted training. Model parameters, momentum buffers, and scheduler state also matched exactly in this deterministic CPU example.

Omitting the scheduler state instead produced 0.05, 0.05, 0.025. The first resumed rate matched, but the following reduction arrived late. Compare a continuation across a schedule boundary when testing resume behavior. Our PyTorch checkpoint guide covers files, random-number state, and broader training-resume requirements.

Check the schedule before a long run

• Unexpectedly fast decay: count scheduler calls and check whether an epoch-based call moved inside the batch loop.
• No plateau reduction: inspect the metric direction, improvement threshold, patience, cooldown, and minimum rate.
• Wrong first rate: inspect the optimizer after constructing the scheduler and check call order.
• Changed resume behavior: compare scheduler state, optimizer groups, and the exact save boundary.

Keep a short record of rates used, completed updates, validation events, and configuration. Compare candidate schedules with a consistent training and validation setup. A correct learning-rate trace establishes that the schedule ran as intended; whether it helps the model still requires an experiment.