🎓 A1Learner
🔍

PyTorch What Is an Optimizer

An optimizer is an object that holds a model's parameters and updates them. Each update moves a parameter against its gradient, and step() applies one update.

The update by hand

A by-hand fit updates each value with its own line of code. Two values need two update lines, and bigger models need far more.

w = w - rate * w.grad
b = b - rate * b.grad

Handing the updates over

An optimizer is built from the model's parameters and a learning rate. The learning rate sets how big each update is.

model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

One step

step() applies one update to every parameter at once. The loss drops after a single step.

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)

loss = loss_fn(model(x), y)
print(f"Before: {loss.item():.4f}")
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss = loss_fn(model(x), y)
print(f"After one step: {loss.item():.4f}")

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)

for step in range(81):
    predictions = model(x)
    loss = loss_fn(predictions, y)
    if step % 20 == 0:
        print(f"Step {step}: loss {loss.item():.4f}")
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

weight = model.weight.item()
bias = model.bias.item()
print(f"Learned weight: {weight:.2f}")
print(f"Learned bias: {bias:.2f}")

Key points

  • Parameter: a value inside a model that training can change, like a weight or a bias
  • Gradient: how much the loss moves when a parameter moves
  • Use model.parameters() to hand a model's parameters to an optimizer
  • Use optimizer.step() to apply one update to every parameter
  • Use optimizer.zero_grad() to clear old gradients before the next backward()
  • The optimizer is the part of training that turns gradients into new parameter values

In this example, the data follows the rule y = 2x + 1. The seed is set before the model is created, so the start is repeatable. A Linear model and an SGD optimizer are built next. Each round computes the predictions and the loss first. The round then clears gradients, runs backward, and lets step() update the model. The loss falls from 28.1172 to 0.0001 in 80 steps. The learned weight and bias print as 1.99 and 1.02, close to the true 2 and 1.

Note

The model computes the loss and the gradients, and the optimizer only applies the updates

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does an optimizer's step() method do?