🎓 A1Learner
🔍

PyTorch Learning Rate Schedules

A learning rate schedule is a rule that changes the learning rate during training. A scheduler object applies the rule once per epoch.

StepLR drops in steps

StepLR holds the rate flat, then multiplies it by gamma. With step_size 3 and gamma 0.5, the rate halves every three epochs.

import torch

torch.manual_seed(0)
model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.5)
rates = []
for epoch in range(7):
    rates.append(f"{optimizer.param_groups[0]['lr']:.4f}")
    optimizer.step()
    scheduler.step()
print(", ".join(rates))

Building a cosine scheduler

CosineAnnealingLR follows a cosine curve, a smooth curve that falls slowly at both ends. T_max sets the length of the fall in epochs.

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=10
)

The step order

The optimizer steps first in each epoch. The scheduler steps after it, once per epoch.

optimizer.step()
scheduler.step()

Example

import torch

torch.manual_seed(0)
x = torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[1.0], [3.0], [5.0], [7.0], [9.0]])
model = torch.nn.Linear(1, 1)
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=10
)

for epoch in range(11):
    rate = optimizer.param_groups[0]["lr"]
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch}: lr {rate:.4f}, loss {loss.item():.4f}")
    scheduler.step()

Key points

  • Scheduler: an object that updates the learning rate on a set plan
  • Use StepLR to drop the rate by a factor every few epochs
  • Use CosineAnnealingLR to ease the rate down along a smooth curve
  • Use scheduler.step() once per epoch, after optimizer.step()
  • T_max: the number of epochs over which the cosine curve falls
  • A schedule replaces one fixed rate with a planned path of rates

In this example, SGD trains a line model under a cosine schedule. The schedule gets T_max 10, so the fall spans ten epochs. Each epoch prints the current rate before the scheduler steps. The rate starts at 0.0500 and eases down to 0.0000 at epoch 10. The fall is gentle at the start and gentle again at the end. The loss falls from 28.1172 to 0.0117 along the way. The last two losses match, because the rate has nearly reached zero.

Note

The scheduler reads the rate from the param groups, and param_groups[0]['lr'] shows the current rate

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

How often does a scheduler usually step?