🎓 A1Learner
🔍

PyTorch Finding a Learning Rate

A learning rate sweep is a set of short training runs at different rates. The final losses show which rate fits the problem.

One shared start

Every run in a fair sweep starts from the same model. Seeding before each creation gives all six runs this same opening loss.

import torch

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]])
loss_fn = torch.nn.MSELoss()
torch.manual_seed(0)
model = torch.nn.Linear(1, 1)
print(f"{loss_fn(model(x), y).item():.4f}")

One short run

A single sweep run trains for a fixed 30 steps. Only the rate changes between runs.

torch.manual_seed(0)
model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=rate)
for step in range(30):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

Reading the table

Small rates leave the loss high after 30 steps. The loss falls as the rate grows, until the rate turns destructive.

for rate in [0.001, 0.003, 0.01, 0.03, 0.1, 0.3]:
    losses.append(train(rate))

Example

import torch

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]])
loss_fn = torch.nn.MSELoss()

for rate in [0.001, 0.003, 0.01, 0.03, 0.1, 0.3]:
    torch.manual_seed(0)
    model = torch.nn.Linear(1, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=rate)
    for step in range(30):
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        final = loss_fn(model(x), y).item()
    print(f"lr {rate}: final loss {final:.4g}")

Key points

  • Sweep: a set of short runs that differ in one setting only
  • Use the same seed in every run so all runs share one start
  • Use a short run for the sweep, and save the long run for the winner
  • The best rate in this sweep is 0.1, with a final loss of 0.0003833
  • Rate 0.3 exploded in this sweep, and its final loss printed as 1.808e+30
  • A sweep turns rate picking from a guess into a small measured table

In this example, six SGD runs fit the line y = 2x + 1 for 30 steps each. The seed resets before each model, so every run starts at loss 28.1172. Rate 0.001 crawls and still ends at 12.52. The final loss falls as the rate rises toward 0.1. Rate 0.1 ends lowest, at 0.0003833. Rate 0.3 overshoots so hard that the loss ends at 1.808e+30. The table picks 0.1 for a longer run on this problem.

Note

The winning rate belongs to this problem and this model, and a new problem needs a new sweep

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

In the sweep, rate 0.3 ended with a loss of 1.808e+30. What does that result show?