M2.5 — Training Loops, Loss Functions and GPU Basics

Learning objectives

After this lesson, you should be able to:

• Name the four pieces of training: model, loss function, optimizer and data.
• Write the five-step training loop in the right order.
• Explain why zero_grad() is needed.
• Describe an epoch and a mini-batch, and read a DataLoader loop.
• Evaluate a model on held-out data with model.eval() and torch.no_grad().
• Move a model and its data to a GPU, and recognise a device mismatch.

The one-sentence idea

Training is a loop that repeats four steps on batches of data: predict, measure the loss, compute gradients and update the parameters.

Predict, score, differentiate, update. Repeat.

The training loop is the engine

M2.4 described one training step: forward pass, loss, backward pass, update. A training loop repeats that step over a dataset until the model is good enough. Every training script for a robot policy, including the ones behind VLA models, is a variation of the same few lines, so learning to read this loop is learning to read all of them.

The examples use a model small enough to follow by eye: it learns the line y = 2x + 1 from data. The structure is identical to a model with a billion parameters. Only the model, the data and the amount of compute change.

The four pieces

Model

The function with parameters, such as an nn.Linear layer or a whole network. It maps inputs to predictions.

Loss function

Compares predictions with targets and returns one number. It says how wrong the model is.

Optimizer

Holds the parameters and updates them from their gradients. It applies the gradient-descent step, or a smarter variant.

Data

Inputs and targets, usually supplied a mini-batch at a time. It is what the model learns from.

Common choices of loss: nn.MSELoss (mean squared error) when the output is a continuous number such as a joint angle, and nn.CrossEntropyLoss when the output is a choice among classes, such as the next token from a vocabulary. Common optimizers are SGD, the plain gradient-descent step of M2.4, and Adam, which adapts the step for each parameter and is the usual default for large models.

A complete training loop

First the data: a line y = 2x + 1 sampled at 20 points. The points are split into a training set, used to update the model, and a validation set that is only used to check it. M6 explains why that separation matters; for now, notice that the model never trains on the validation points.

import torch
from torch import nn

x = torch.linspace(-1, 1, 20).unsqueeze(1)    # shape (20, 1)
y = 2 * x + 1
x_train, y_train = x[::2], y[::2]             # every other point
x_val,   y_val   = x[1::2], y[1::2]           # the points in between

model = nn.Linear(1, 1)
with torch.no_grad():                          # start from known values
    model.weight.fill_(0.0)
    model.bias.fill_(0.0)

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
print(x_train.shape, x_val.shape)

torch.Size([10, 1]) torch.Size([10, 1])

Now the loop itself. Each pass through it does the five steps that matter, in this order:

for step in range(200):
    pred = model(x_train)                # 1. forward pass
    loss = loss_fn(pred, y_train)        # 2. compute the loss
    optimizer.zero_grad()                # 3. clear old gradients
    loss.backward()                      # 4. backward pass: fill in the gradients
    optimizer.step()                     # 5. update the parameters
    if step in (0, 50, 199):
        print(step, round(loss.item(), 4))

print(round(model.weight.item(), 2), round(model.bias.item(), 2))

0 2.2632
50 0.0008
199 0.0
2.0 1.0

The loss starts at 2.26 and falls toward zero, and the model recovers the slope 2 and offset 1 from the data. One full pass through all the training data is called an epoch. Here each step used the whole training set, so one step was one epoch.

model.eval()                              # switch layers like dropout to evaluation behaviour
with torch.no_grad():                     # no gradients needed, saves memory and time
    val_loss = loss_fn(model(x_val), y_val)
print(val_loss.item() < 1e-4)

True

torch.no_grad() turns off gradient tracking when you only want predictions. model.eval() puts the model in evaluation mode, which matters for layers whose behaviour differs between training and testing; call model.train() to switch back. A low loss on data the model never trained on is the evidence that it learned the pattern and did not just memorise.

Mini-batches

Real datasets are too large to process at once, so each step uses a mini-batch, a small group of examples. A DataLoader wraps a dataset and hands out batches, optionally in a random order. A nested loop then runs over epochs and, inside each, over batches. Using batches also gives updates that are noisy but cheap, and that noise often helps training.

from torch.utils.data import TensorDataset, DataLoader

loader = DataLoader(TensorDataset(x_train, y_train), batch_size=4, shuffle=False)
for xb, yb in loader:
    print(xb.shape, yb.shape)

torch.Size([4, 1]) torch.Size([4, 1])
torch.Size([4, 1]) torch.Size([4, 1])
torch.Size([2, 1]) torch.Size([2, 1])

The 10 training points arrive as batches of 4, 4 and 2, and the leading dimension is the batch dimension from M2.3. In a real run you would use shuffle=True for training, and put the five loop steps inside for xb, yb in loader:.

The bugs everyone meets once

Forgetting zero_grad(). PyTorch adds new gradients to existing ones. Without clearing them, each step uses the sum of all earlier gradients:

w = torch.tensor(1.0, requires_grad=True)
(w * 3).backward()
(w * 3).backward()
print(w.grad)          # 3 + 3
w.grad.zero_()
print(w.grad)

tensor(6.)
tensor(0.)

Other common mistakes: evaluating without model.eval() and torch.no_grad(); validating on data the model trained on, which hides overfitting; a learning rate so large that the loss grows or becomes nan; and a mismatch between prediction and target shapes, which can silently broadcast to the wrong size. When a loss behaves oddly, print the shapes and a few values first.

GPU basics

To train on a GPU, the model and every tensor it touches must be on the same device. Choose the device once, move the model once, and move each batch inside the loop:

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)                   # move the parameters once

for xb, yb in loader:
    xb, yb = xb.to(device), yb.to(device)  # move each batch
    loss = loss_fn(model(xb), yb)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

This is not run here because its behaviour depends on your hardware. A GPU is much faster for large models, but it has limited memory, and the forward pass has to keep its intermediate values for the backward pass (as M2.4 described). If a run fails with an out-of-memory error, the usual remedies are a smaller batch size or a smaller model. Printing a loss value with loss.item() copies it back to the CPU, so doing it every step on a GPU can slow a run down.

For this prerequisite material you do not need a GPU. The same code runs on a CPU, only more slowly for big models.

Common misconceptions

Engineering takeaways

1. Training needs a model, a loss function, an optimizer and data.
2. The loop is: forward pass, loss, zero_grad(), backward(), step().
3. An epoch is one pass over the training data. A mini-batch is a small group of examples per step.
4. MSE is the usual loss for continuous outputs, and cross-entropy for choices among classes.
5. Adam is a common default optimizer, and SGD is the plain gradient-descent step.
6. Evaluate with model.eval() and torch.no_grad(), on data the model did not train on.
7. Gradients accumulate unless cleared.
8. On a GPU, the model and every batch must be on the same device.

Read any training script by finding these five lines, and the rest is detail.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

That completes the Python and deep-learning prerequisites. The next module covers the architecture that almost every modern VLA is built on.

Back to the VLA roadmap prerequisites