🎓 A1Learner
🔍
0 of 79 lessons completed

PyTorch The Learning Rate

The learning rate is a number that sets how big each optimizer step is. A rate that fits the problem makes the loss fall steadily.

The rate lives in the optimizer

The learning rate is set once, when the optimizer is created. It can be read back from the param groups.

import torch

torch.manual_seed(0)
model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
print(optimizer.param_groups[0]["lr"])

A tiny rate moves a little

One step at rate 0.001 moves the weight only a small amount.

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

print(f"Weight before: {model.weight.item():.4f}")
loss = loss_fn(model(x), y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Weight after: {model.weight.item():.4f}")

A huge rate breaks the loss

Twenty steps at rate 10 wreck the fit completely. The loss ends as nan, which means not a number.

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=10.0)

for step in range(20):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

loss = loss_fn(model(x), y)
print(f"Loss after 20 steps: {loss.item():.4f}")

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()


def train(rate):
    torch.manual_seed(0)
    model = torch.nn.Linear(1, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=rate)
    losses = []
    for step in range(31):
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        if step % 10 == 0:
            losses.append(f"{loss.item():.4f}")
        loss.backward()
        optimizer.step()
    return losses


for rate in [0.001, 0.05, 10.0]:
    losses = train(rate)
    print(f"Rate {rate}: {', '.join(losses)}")

Key points

  • Learning rate: the number that scales every step an optimizer takes
  • Use the lr argument to set the rate when the optimizer is created
  • A tiny rate makes training crawl, and the loss falls very slowly
  • A huge rate can overshoot so badly that the loss becomes inf or nan
  • Inf: the marker for a number too big to store
  • Nan: a special value that means not a number
  • The right rate sits between crawling and exploding, and a quick test shows where it is

In this example, one train function fits the line y = 2x + 1. The function seeds the model at the start, so every rate gets the same start. Rate 0.001 crawls, and the loss is still 12.5209 after 30 steps. Rate 0.05 falls fast and ends near zero. Rate 10.0 explodes, and the loss passes inf and ends as nan. The middle rate wins this test.

Note

There is no single best learning rate, and a new problem usually needs a new short test

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A training loss turns into nan after a few steps. Which change is the first thing to try?