PyTorch Training Loop

The training loop is the heart of every PyTorch project. It repeats the same sequence of steps — predict, measure error, compute gradients, update weights — thousands of times until the model learns from the data. Every deep learning model, from a tiny classifier to a billion-parameter language model, trains using this same loop.

The Five Steps in Every Training Iteration

┌─────────────────────────────────────────────────────┐
│  TRAINING LOOP STEPS (repeat for each batch)        │
│                                                     │
│  1. optimizer.zero_grad()  ← clear old gradients    │
│  2. predictions = model(X) ← forward pass           │
│  3. loss = loss_fn(pred,y) ← compute error          │
│  4. loss.backward()        ← compute gradients      │
│  5. optimizer.step()       ← update weights         │
└─────────────────────────────────────────────────────┘

A Complete Training Loop

import torch
import torch.nn as nn
import torch.optim as optim

# Setup
model = nn.Sequential(
    nn.Linear(4, 16),
    nn.ReLU(),
    nn.Linear(16, 1)
)

optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss()

# Fake dataset: 100 samples, 4 features each
X = torch.rand(100, 4)
y = torch.rand(100, 1)

# Training loop
num_epochs = 50

for epoch in range(num_epochs):
    model.train()                          # set training mode

    optimizer.zero_grad()                  # step 1: clear gradients
    predictions = model(X)                 # step 2: forward pass
    loss = loss_fn(predictions, y)         # step 3: compute loss
    loss.backward()                        # step 4: compute gradients
    optimizer.step()                       # step 5: update weights

    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{num_epochs}]  Loss: {loss.item():.4f}")

What Each Step Does

Step 1 – zero_grad()

PyTorch accumulates gradients by default. If you forget to clear them, each backward pass adds to the gradients from the previous step, causing incorrect weight updates. Always call zero_grad() first.

Step 2 – Forward Pass

The model processes the input batch and produces predictions. PyTorch records every operation in the computation graph during this step.

Step 3 – Loss Computation

The loss function compares predictions to true labels and returns a single scalar value representing total error for this batch.

Step 4 – backward()

PyTorch walks the computation graph backwards and computes the gradient of the loss with respect to every learnable parameter.

Step 5 – optimizer.step()

The optimizer reads the gradients and adjusts each weight. Parameters with large gradients (that contributed heavily to the error) change more than parameters with small gradients.

Adding a Validation Loop

Train the model on training data, then evaluate it on validation data after each epoch. Validation shows whether the model is generalizing or merely memorizing training examples.

X_train, y_train = torch.rand(80, 4), torch.rand(80, 1)
X_val,   y_val   = torch.rand(20, 4), torch.rand(20, 1)

for epoch in range(num_epochs):
    # Training phase
    model.train()
    optimizer.zero_grad()
    pred = model(X_train)
    loss = loss_fn(pred, y_train)
    loss.backward()
    optimizer.step()

    # Validation phase
    model.eval()
    with torch.no_grad():
        val_pred = model(X_val)
        val_loss = loss_fn(val_pred, y_val)

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}  Train Loss: {loss.item():.4f}  Val Loss: {val_loss.item():.4f}")

Training vs Validation Loss Pattern

Healthy training:
  Epoch  Train Loss   Val Loss
    10     0.50        0.52
    20     0.30        0.33
    50     0.10        0.13
  → Both losses decrease and stay close together ✓

Overfitting:
  Epoch  Train Loss   Val Loss
    10     0.50        0.52
    20     0.10        0.45
    50     0.01        0.60
  → Train loss keeps falling, val loss rises ✗

Training with Mini-Batches (DataLoader)

In practice, you don't pass all data at once. You use a DataLoader to split data into batches and loop over them.

from torch.utils.data import TensorDataset, DataLoader

dataset = TensorDataset(X, y)
loader  = DataLoader(dataset, batch_size=16, shuffle=True)

for epoch in range(num_epochs):
    model.train()
    total_loss = 0.0

    for batch_X, batch_y in loader:       # loop over batches
        optimizer.zero_grad()
        pred = model(batch_X)
        loss = loss_fn(pred, batch_y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(loader)
    print(f"Epoch {epoch+1}  Avg Loss: {avg_loss:.4f}")

Loop Diagram with DataLoader

Dataset: 1000 samples, batch_size=32 → 32 batches per epoch

Epoch 1:
  Batch 1  [32 samples] → step 1-5 → weight update
  Batch 2  [32 samples] → step 1-5 → weight update
  ...
  Batch 32 [32 samples] → step 1-5 → weight update
  ↓ compute avg loss for epoch 1

Epoch 2:
  (data reshuffled)
  Batch 1  [different 32 samples] → ...
  ...

Summary

The training loop repeats five steps: clear gradients, run forward pass, compute loss, call backward, call optimizer.step. Always call model.train() before the training loop and model.eval() before validation. Wrap validation code in torch.no_grad() to skip gradient computation. Use a DataLoader to process data in mini-batches. Monitor both training and validation loss to detect overfitting. The training loop is the same structure regardless of model complexity — only the contents of the model and loss function change.