PyTorch Anatomy of a Training Loop

A training loop is the repeated cycle that trains a model one step at a time. Each step runs the same five lines in a fixed order.

The five lines

One training step is these five lines. The order is part of the method.

pred = model(x)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

Swap two lines and training stops

Here zero_grad() runs after backward(). It wipes the fresh gradients, so step() has nothing to move.

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

The right order learns

The same setup in the right order drops the loss at every step. The loop is the five lines again and again.

for step in range(3):
    pred = model(x)
    loss = loss_fn(pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print("Step", step, "loss:", round(loss.item(), 4))

Example

import torch
from torch import nn


def train(model, x, y, steps, lr):
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    for step in range(steps):
        pred = model(x)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return loss.item()


torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]])
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()

start = loss_fn(model(x), y).item()
final = train(model, x, y, steps=200, lr=0.05)

print("Start loss:", round(start, 4))
print("Final loss:", round(final, 4))
print("Weight:", round(model.weight.item(), 2))
print("Bias:", round(model.bias.item(), 2))

Key points

Forward pass: the model turns inputs into predictions
Loss: one number that scores the predictions against the targets
Use optimizer.zero_grad() to clear old gradients before new ones are made
Use loss.backward() to fill each parameter's gradient from the loss
Use optimizer.step() to move each parameter against its gradient
One step is the five lines once, in this order

In this example, the train function holds the five lines. It builds a loss and an optimizer, then loops the steps. Each step runs forward, loss, zero_grad, backward, and step. The main code sets seed 0 before the model is built. The seed makes the starting weight repeatable. The data follows the line y = 2x + 1 exactly. The start loss prints before any training. The function then trains for 200 steps at rate 0.05. The final loss rounds to 0.0 at four decimals. The learned weight and bias match the true line. The next lesson trains on noisy data, where the loss cannot reach zero.

Note

Gradients add up by default, so a missed zero_grad() trains on stale gradients

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn


def train(model, x, y, steps, lr):
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    for step in range(steps):
        pred = model(x)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return loss.item()


torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]])
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()

start = loss_fn(model(x), y).item()
final = train(model, x, y, steps=200, lr=0.05)

print("Start loss:", round(start, 4))
print("Final loss:", round(final, 4))
print("Weight:", round(model.weight.item(), 2))
print("Bias:", round(model.bias.item(), 2))

Exercise

In one training step, what must run right before loss.backward()?

optimizer.step() optimizer.zero_grad() model.eval() torch.save()