🎓 A1Learner
🔍
0 of 87 lessons completed

PyTorch Warmup and Decay

Warmup is a short opening phase where the learning rate rises from near zero. A decay phase then lowers the rate for the rest of training.

The shape as factors

The factor climbs by equal steps for five epochs. The factor then falls slowly toward zero.

def rate_for_epoch(epoch):
    if epoch < 5:
        return (epoch + 1) / 5
    return 1.0 - (epoch - 5) / 15


print([round(rate_for_epoch(epoch), 2) for epoch in range(7)])

Wiring LambdaLR

LambdaLR takes the optimizer and the factor function. The base rate here is the optimizer's rate, 0.1.

scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer, lr_lambda=rate_for_epoch
)

Warmup rates on an optimizer

The printed rates are the factor times the base rate 0.1. Epoch 5 and epoch 6 both sit at the peak.

import torch


def rate_for_epoch(epoch):
    if epoch < 5:
        return (epoch + 1) / 5
    return 1.0 - (epoch - 5) / 15


torch.manual_seed(0)
model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=rate_for_epoch)
rates = []
for epoch in range(6):
    rates.append(f"{optimizer.param_groups[0]['lr']:.4f}")
    optimizer.step()
    scheduler.step()
print(", ".join(rates))

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.1)


def rate_for_epoch(epoch):
    if epoch < 5:
        return (epoch + 1) / 5
    return 1.0 - (epoch - 5) / 15


scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=rate_for_epoch)

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

Key points

  • Warmup: the opening epochs where the rate climbs to its peak
  • Decay: the later phase where the rate falls toward zero
  • Use LambdaLR with a function that maps an epoch number to a rate factor
  • The scheduler multiplies the base rate by the returned factor
  • Warmup plus decay gives the rate a rise, a peak, and a fall

In this example, LambdaLR shapes the rate for a line model over 20 epochs. The factor function rises for the first five epochs. The printed rate climbs from 0.0200 to the 0.1000 peak at epoch 5. The rate holds at the peak for one more epoch. The factor then decays in equal steps for the rest of the run. The final printed rate is 0.0067 at epoch 20. The loss falls from 28.1172 to 0.0041 along the table.

Note

The epoch number inside the LambdaLR function starts at 0

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does warmup start with a tiny learning rate?