# The canonical training loop — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-training-loop

> Write the five-line training step in the correct order, track loss safely, and sanity-check the starting loss.

## Five lines, in this order

Almost every PyTorch training loop has an outer loop over **epochs** and an inner loop over **batches**. Inside, the order is fixed:

1. `optimizer.zero_grad()` clears old gradients, because they accumulate.
2. `logits = model(x)` runs the forward pass and builds the graph.
3. `loss = loss_fn(logits, y)` reduces everything to one scalar.
4. `loss.backward()` fills `.grad` for every parameter.
5. `optimizer.step()` updates the parameters using those gradients.

Before these, move the batch to the same device as the model. Call `model.train()` at the start of each epoch, because evaluation code elsewhere may have switched the model to eval mode.

For logging, accumulate `loss.item()` (a Python float), never the `loss` tensor, which would keep every batch's graph in memory. `zero_grad()` may also go right after `step()`; what matters is that it happens between one `step()` and the next `backward()`. Swapping `backward()` and `step()` means each update uses the previous batch's gradients, or none.

## Tracing one iteration

Follow one batch through the loop. The state shows the model mode, whether a graph exists, what `.grad` holds and whether the weights have changed.

```python
for epoch in range(num_epochs):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        logits = model(x)
        loss = loss_fn(logits, y)
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch} last batch loss {loss.item():.4f}")
```

- Line 1: The outer loop counts epochs: full passes over the training data.
- Line 2: Switch dropout and batch norm to training behaviour. Validation at the end of the previous epoch may have left the model in eval mode.
- Line 3: The DataLoader yields one shuffled mini-batch of inputs and labels.
- Line 4: Data must be on the same device as the model's parameters, or the forward pass raises a device mismatch error.
- Line 5: Old gradients are cleared first, because backward adds to whatever is already in `.grad`.
- Line 6: The forward pass computes logits and autograd records every operation.
- Line 7: The loss reduces the batch to one scalar, the root of the graph.
- Line 8: Backward walks the graph in reverse, fills `.grad` for every parameter, then frees the graph.
- Line 9: Only now do the weights change, using this batch's gradients.
- Line 10: After all batches, the epoch summary prints. `.item()` gives a plain float, so no graph is kept alive.

## A complete, self-contained example

Synthetic 2-D points labelled by which side of a line they fall on. Full-batch training keeps it short; the structure is the same as with a DataLoader.

```python
import torch
from torch import nn

torch.manual_seed(0)
X = torch.randn(256, 2)
y = (X[:, 0] + X[:, 1] > 0).long()        # label: above or below the line

model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

for epoch in range(50):
    optimizer.zero_grad()
    loss = loss_fn(model(X), y)
    loss.backward()
    optimizer.step()

with torch.no_grad():
    acc = (model(X).argmax(dim=1) == y).float().mean().item()
print(f"final loss {loss.item():.3f}, train accuracy {acc:.2%}")
```

## Check the very first loss

With random initial weights, a C-class classifier predicts roughly uniform probabilities, so the first cross-entropy loss should be close to ln(C): about 2.30 for 10 classes, 0.69 for 2. A first loss far from that points to a bug in labels, the loss choice or the final layer.

**Quiz:** Which order of calls inside the batch loop is correct?

- [ ] `backward()`, `zero_grad()`, forward and loss, `step()`
- [x] `zero_grad()`, forward and loss, `backward()`, `step()`
- [ ] forward and loss, `step()`, `backward()`, `zero_grad()`
- [ ] `step()`, forward and loss, `zero_grad()`, `backward()`

*Answer:* `zero_grad()`, forward and loss, `backward()`, `step()`. Clear gradients, compute the loss, back-propagate, then update. Any other order either uses stale gradients or wipes the new ones before they are used.