๐ŸŽ“ A1Learner
๐Ÿ”

PyTorch Plateau and OneCycle Schedulers

ReduceLROnPlateau is a scheduler that lowers the rate after a watched score stalls. OneCycleLR is a scheduler that raises the rate once and then lowers it.

Plateau waits, then halves

Patience of one means two stalled scores trigger the drop. Three flat metrics take the rate from 0.1 to 0.05.

import torch

model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.10)
plateau = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.5, patience=1)
for metric in [1.0, 1.0, 1.0]:
    plateau.step(metric)
print(plateau.get_last_lr()[0])

OneCycle climbs first

The rate starts far below the peak and climbs batch by batch. Three batches reach the peak of 0.1 in this short run.

import torch

model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
onecycle = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=0.10, total_steps=10)
for step in range(3):
    print(f"{onecycle.get_last_lr()[0]:.4f}")
    optimizer.step()
    onecycle.step()

Stepping OneCycle per epoch strands the rate

A ten-batch plan stepped only three times ends at 0.0950, just past the peak. The long fall of the arc never runs.

import torch

model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
onecycle = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=0.10, total_steps=10)
for epoch in range(3):
    optimizer.step()
    onecycle.step()
print(f"{onecycle.get_last_lr()[0]:.4f}")

Example

import torch

torch.manual_seed(0)
model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.10)
plateau = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.5, patience=1)
metrics = [1.00, 0.90, 0.90, 0.90, 0.85, 0.85, 0.85]
print("Plateau trace:")
for metric in metrics:
    plateau.step(metric)
    print(f"  metric {metric:.2f} -> lr {plateau.get_last_lr()[0]:.4f}")

torch.manual_seed(0)
model2 = torch.nn.Linear(2, 1)
optimizer2 = torch.optim.SGD(model2.parameters(), lr=0.01)
onecycle = torch.optim.lr_scheduler.OneCycleLR(
    optimizer2, max_lr=0.10, total_steps=10)
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
y = torch.tensor([[1.0], [2.0]])
loss_fn = torch.nn.MSELoss()
print("OneCycle trace:")
for step in range(10):
    lr = onecycle.get_last_lr()[0]
    optimizer2.zero_grad()
    loss = loss_fn(model2(x), y)
    loss.backward()
    optimizer2.step()
    onecycle.step()
    print(f"  batch {step + 1} -> lr {lr:.4f}")

Key points

  • Metric: the score a scheduler watches, often the validation loss
  • Patience: how many stalled epochs pass before the rate drops
  • Use plateau.step(metric) once per epoch, with that epoch's metric
  • Use onecycle.step() once per batch, and never once per epoch
  • Total steps: every batch in the whole run, counted before training starts
  • Plateau reacts to evidence, and OneCycle follows one fixed arc

In this example, the Plateau half feeds a scripted metric list to a fresh scheduler. The metric stalls twice at 0.90, and the rate halves to 0.0500. It stalls twice more at 0.85, and the rate halves again to 0.0250. The OneCycle half then trains a tiny model for ten batches. Its rate climbs from 0.0040 to the 0.1000 peak at batch 3. The rate then falls batch by batch and ends at 0.0000.

Note

Use get_last_lr() to read the current rate from either scheduler

Try it Yourself

Edit the code, then press Run.

Running is simulated ยท Python Playground

Exercise

When does ReduceLROnPlateau lower the learning rate?