Lesson 4

Loss functions and the full training loop

Introduction

Inputs 0, 1, 2, and 3 should map to targets 1, 3, 5, and 7. A training loop repeatedly measures the mismatch, computes gradients, updates parameters, and clears old gradients.

Learning goal

Implement and explain the five actions in a complete training step while keeping training separate from evaluation.

Before you start

Tensors, autograd, modules, mean squared error, and a basic Python loop.

Lesson plan

  1. Name loss, optimizer, epoch, batch, and the purpose of every loop component.
  2. Train a straight-line model while tracing predictions, loss, gradients, and updates.
  3. Switch to evaluation correctly and modify the example without leaking gradients.

Why a loop is needed

A freshly created model usually has random parameters, so its predictions are poor. One training step measures the current error and makes a small correction. Thousands of steps can gradually produce useful behavior. The model architecture may change from a line to a transformer, but the central sequence stays almost the same.

for x, target in data:
    optimizer.zero_grad()
    prediction = model(x)
    loss = loss_function(prediction, target)
    loss.backward()
    optimizer.step()

The order matters. First clear gradients left by the previous step. Next run the model forward. Then turn its prediction into one loss value. backward() computes gradients. Finally, step() changes parameters according to the optimizer rule.

Terms you will see

BatchA small group of examples processed together.
EpochOne pass through the complete training set.
Loss functionA rule that turns predictions and targets into an error score.
OptimizerAn algorithm that uses gradients to update parameters.
Learning rateThe update step size. Too large may be unstable; too small may be slow.
ValidationMeasurement on examples that are not used for updates.

Worked example: learn a straight line

We will learn the relationship target = 2x + 1. Both x and target have shape (4, 1): four examples with one feature each. nn.Linear(1, 1) holds one weight and one bias. Mean squared error averages (prediction - target)² across the batch. SGD means stochastic gradient descent; here it updates parameters by subtracting learning rate times gradient.

Predict first: if training learns the rule, what should it return for x = 4? Write the answer before running the loop. During the first few steps, also expect the parameters to change only after optimizer.step().

import torch
from torch import nn

torch.manual_seed(0)
x = torch.tensor([[0.0], [1.0], [2.0], [3.0]])
target = torch.tensor([[1.0], [3.0], [5.0], [7.0]])

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

for step in range(300):
    optimizer.zero_grad()
    prediction = model(x)
    loss = loss_function(prediction, target)
    loss.backward()
    optimizer.step()

weight = model.weight.item()
bias = model.bias.item()
with torch.no_grad():
    answer = model(torch.tensor([[4.0]])).item()

print(round(weight, 2), round(bias, 2))
print(round(answer, 2))

# Expected, up to tiny floating-point differences:
# 2.0 1.0
# 9.0

At each step, the forward pass uses the current weight and bias. The loss connects the prediction to the target. The backward pass fills model.weight.grad and model.bias.grad. The optimizer then updates those two parameters. Autograd alone would leave them unchanged.

The expected answer is 9. If the loss falls while the answer stays near its initial value, print the parameters before and after step(). Unchanged parameters usually mean the optimizer does not own them, step() was skipped, or gradients were disabled.

Training and evaluation are separate jobs

Training uses model.train(), gradient tracking, and optimizer updates. Validation uses model.eval() and no parameter update. Some layers, especially dropout and batch normalization, behave differently in training and evaluation modes. Evaluation also does not need a backward graph, so torch.no_grad() reduces memory use.

model.eval()
with torch.no_grad():
    validation_prediction = model(x)
    validation_loss = loss_function(validation_prediction, target)
print(round(validation_loss.item(), 6))  # Expected: near 0.0

model.eval() does not turn off gradients by itself. Likewise, torch.no_grad() does not place dropout into evaluation behavior. For ordinary validation, use both. Call model.train() again before resuming training.

Common pitfalls

  • Forgetting zero_grad() adds new gradients to old ones.
  • Calling step() before backward() gives the optimizer no new gradient to use.
  • A falling training loss does not prove good generalization; also measure a validation set.
  • Use compatible prediction and target shapes. Broadcasting can hide a shape mistake.
  • Very large learning rates can make the loss jump or become nan.

Try it

Change the relationship to target = -3x + 2. Which two learned values should you expect? Keep the input shape (4, 1) and train a fresh model.

Reveal the worked answer 
import torch
from torch import nn

torch.manual_seed(1)
x = torch.tensor([[0.0], [1.0], [2.0], [3.0]])
y = -3 * x + 2
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

for _ in range(300):
    optimizer.zero_grad()
    loss = nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()

print(round(model.weight.item(), 2))  # Expected: -3.0
print(round(model.bias.item(), 2))    # Expected: 2.0

Recap

A training step clears gradients, predicts, measures loss, computes new gradients, and updates parameters. A batch groups examples; an epoch covers the training set. The learning rate controls update size. Keep validation separate, and use evaluation mode plus disabled gradient recording when no update is needed.

Reference: PyTorch optimization tutorial.