Training Loops, Batches & Epochs

By the end of this lesson you will be able to write a proper training loop: shuffle the data, cut it into mini-batches, update the weights after each one, and track training and validation loss every epoch. You will also understand why almost all deep learning uses mini-batches.

In the last lesson the network looked at all the training data before every update. That works for 300 points. It does not work for 60,000 images, let alone millions. Here is how real training loops handle it.

Epochs, batches and iterations

Three words describe the shape of training:

Picture a teacher in Lucknow correcting 120 answer sheets. She could read all 120 and only then change how she teaches (full batch). She could change her plan after every single sheet (very jumpy). Or she could read a bundle of 20, adjust, then read the next 20. That middle path is the mini-batch, and it is what nearly everyone uses.

There are three styles:

• Full-batch gradient descent: one update per epoch using all examples. Accurate gradient, but slow and memory hungry.
• Stochastic gradient descent (SGD): one update per single example. Very noisy. "Stochastic" means random.
• Mini-batch gradient descent: one update per small batch, typically 16 to 512 examples. In practice people call this "SGD" too.

Step 1: shuffle and cut into mini-batches

We shuffle the order at the start of every epoch. If the data were sorted (all the 0 digits first, then all the 1s), each batch would contain only one class and training would swing back and forth.

import numpy as np

def minibatches(X, y, batch_size, rng):
    order = rng.permutation(len(X))           # a new random order each epoch
    for start in range(0, len(X), batch_size):
        idx = order[start:start + batch_size]
        yield X[idx], y[idx]

X = np.arange(10).reshape(10, 1) * 1.0         # 10 tiny examples: 0..9
y = np.arange(10)
rng = np.random.default_rng(0)

for epoch in range(1, 3):
    print(f"epoch {epoch}")
    for step, (xb, yb) in enumerate(minibatches(X, y, batch_size=4, rng=rng), start=1):
        print(f"  step {step}: labels {yb}, batch shape {xb.shape}")

epoch 1
  step 1: labels [4 6 2 7], batch shape (4, 1)
  step 2: labels [3 5 9 0], batch shape (4, 1)
  step 3: labels [8 1], batch shape (2, 1)
epoch 2
  step 1: labels [2 9 3 6], batch shape (4, 1)
  step 2: labels [0 4 8 7], batch shape (4, 1)
  step 3: labels [5 1], batch shape (2, 1)

Each epoch visits every example exactly once, in a new order. With 10 examples and a batch size of 4, there are 3 steps, and the last batch has only 2 examples. That is fine: the loss is a mean, so a smaller batch still gives a sensible gradient. The number of steps per epoch is the total divided by the batch size, rounded up.

Step 2: mini-batch versus full-batch training

Let us train the digits network both ways for the same 20 epochs. To keep the code short, the network lives in two functions: one forward-and-backward function and one update.

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
X = X / 16.0
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.25, random_state=42, stratify=y)

def init(seed=42):
    rng = np.random.default_rng(seed)
    return {"W1": rng.normal(0, 1, (64, 64)) * np.sqrt(2 / 64), "b1": np.zeros(64),
            "W2": rng.normal(0, 1, (64, 10)) * np.sqrt(2 / 64), "b2": np.zeros(10)}

def forward_backward(p, X, y):
    Z1 = X @ p["W1"] + p["b1"]
    A1 = np.maximum(0, Z1)
    Z2 = A1 @ p["W2"] + p["b2"]
    P = np.exp(Z2 - Z2.max(axis=1, keepdims=True))
    P /= P.sum(axis=1, keepdims=True)
    loss = -np.log(P[np.arange(len(y)), y]).mean()
    dZ2 = P.copy()
    dZ2[np.arange(len(y)), y] -= 1
    dZ2 /= len(y)
    dZ1 = (dZ2 @ p["W2"].T) * (Z1 > 0)
    grads = {"W1": X.T @ dZ1, "b1": dZ1.sum(axis=0), "W2": A1.T @ dZ2, "b2": dZ2.sum(axis=0)}
    return loss, P, grads

def train(batch_size, epochs=20, lr=0.5):
    p, rng, updates = init(), np.random.default_rng(0), 0
    for epoch in range(epochs):
        order = rng.permutation(len(X_train))
        for start in range(0, len(X_train), batch_size):
            idx = order[start:start + batch_size]
            _, _, grads = forward_backward(p, X_train[idx], y_train[idx])
            for k in p:
                p[k] -= lr * grads[k]
            updates += 1
    val_loss, P, _ = forward_backward(p, X_val, y_val)
    return updates, val_loss, np.mean(P.argmax(axis=1) == y_val)

for bs in [len(X_train), 32]:
    updates, val_loss, val_acc = train(bs)
    name = "full batch" if bs == len(X_train) else f"batch {bs}"
    print(f"{name:<10}: {updates:>4} updates, validation loss {val_loss:.3f}, accuracy {val_acc:.3f}")

full batch:   20 updates, validation loss 0.597, accuracy 0.889
batch 32  :  860 updates, validation loss 0.095, accuracy 0.969

Both runs saw every image 20 times. The full-batch run made only 20 updates, so it reached just 89% with a high loss. The mini-batch run made one update per 32 images, over 800 updates in total, and is much better. Each mini-batch gradient is a little noisy, but many decent steps beat a few perfect ones.

Step 3: watch the loss during training

A good training loop records two things each epoch: the training loss (on the data it learns from) and the validation loss (on held-back data it never learns from). Validation data is a separate slice used to check progress while you train; the final test set is kept for the very end.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
X = X / 16.0
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.25, random_state=42, stratify=y)

rng = np.random.default_rng(42)
p = {"W1": rng.normal(0, 1, (64, 64)) * np.sqrt(2 / 64), "b1": np.zeros(64),
     "W2": rng.normal(0, 1, (64, 10)) * np.sqrt(2 / 64), "b2": np.zeros(10)}

def forward_backward(X, y):
    Z1 = X @ p["W1"] + p["b1"]
    A1 = np.maximum(0, Z1)
    Z2 = A1 @ p["W2"] + p["b2"]
    P = np.exp(Z2 - Z2.max(axis=1, keepdims=True))
    P /= P.sum(axis=1, keepdims=True)
    loss = -np.log(P[np.arange(len(y)), y]).mean()
    dZ2 = P.copy()
    dZ2[np.arange(len(y)), y] -= 1
    dZ2 /= len(y)
    dZ1 = (dZ2 @ p["W2"].T) * (Z1 > 0)
    return loss, {"W1": X.T @ dZ1, "b1": dZ1.sum(axis=0), "W2": A1.T @ dZ2, "b2": dZ2.sum(axis=0)}

step_losses, train_hist, val_hist = [], [], []
for epoch in range(1, 16):
    order = rng.permutation(len(X_train))
    epoch_losses = []
    for start in range(0, len(X_train), 32):
        idx = order[start:start + 32]
        loss, grads = forward_backward(X_train[idx], y_train[idx])
        for k in p:
            p[k] -= 0.2 * grads[k]
        step_losses.append(loss)
        epoch_losses.append(loss)
    train_hist.append(np.mean(epoch_losses))
    val_hist.append(forward_backward(X_val, y_val)[0])
    if epoch in (1, 5, 10, 15):
        print(f"epoch {epoch:>2}: train loss {train_hist[-1]:.3f}, validation loss {val_hist[-1]:.3f}")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
ax1.plot(step_losses, linewidth=0.7)
ax1.set_title("Loss of every mini-batch (noisy)")
ax1.set_xlabel("step")
ax2.plot(range(1, 16), train_hist, label="train")
ax2.plot(range(1, 16), val_hist, label="validation")
ax2.set_title("Average loss per epoch")
ax2.set_xlabel("epoch")
ax2.legend()
plt.show()

epoch  1: train loss 1.367, validation loss 0.890
epoch  5: train loss 0.208, validation loss 0.767
epoch 10: train loss 0.112, validation loss 0.164
epoch 15: train loss 0.075, validation loss 0.139

The left chart jumps up and down, because each batch holds different images. The right chart, averaged per epoch, falls steadily, although the validation loss can still wobble from one epoch to the next (look at epoch 5). Always judge training by averages and by the validation curve, never by one noisy batch. When the validation loss stops falling while the training loss keeps going down, the network has started to memorise. You will handle that in Overfitting & Regularisation.

Choosing a batch size

There is no single correct batch size, but these rules of thumb work well:

1. Start with 32 or 64 for small models, and larger (128 to 512) for images on a GPU.
2. Use the largest batch that fits comfortably in memory, then tune the learning rate for it.
3. If you make the batch much larger, you can usually raise the learning rate too, because each gradient is less noisy.
4. Very large batches give fewer updates per epoch and can generalise slightly worse; very small batches are slow because the hardware is underused.

Powers of two (32, 64, 128) are a habit from GPU memory layouts. They are convenient, not required.

The same loop in PyTorch

PyTorch's DataLoader does the shuffling and batching for you. The loop has the same shape as ours. This is plain code; you will build it fully in A Full PyTorch Training Loop.

for epoch in range(epochs):
    model.train()
    for xb, yb in train_loader:          # DataLoader(..., batch_size=32, shuffle=True)
        loss = loss_fn(model(xb), yb)    # forward + loss
        optimizer.zero_grad()            # clear old gradients
        loss.backward()                  # backward
        optimizer.step()                 # update

Common mistakes

• Not shuffling. Sorted data gives batches of one class and the loss swings wildly.
• Shuffling the validation or test set into training. Shuffle only within the training set, after the split.
• Confusing epochs with steps. "Trained for 10" means little until you say which. Learning rate schedules and logs need to be clear.
• Reading too much into one batch. A single batch loss is noisy; compare epoch averages.

Key points

• An epoch is one full pass over the training data; an iteration is one update on one batch.
• Mini-batches give many cheap, slightly noisy updates and usually train much faster than full batches.
• Shuffle the training data at the start of every epoch.
• Track average training loss and validation loss every epoch.
• Batch sizes of 32 to 512 are common; tune the learning rate together with the batch size.

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
X = X / 16.0
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.25, random_state=42, stratify=y)

def run(batch_size, epochs=20, lr=0.5):
    rng = np.random.default_rng(42)
    W1, b1 = rng.normal(0, 1, (64, 64)) * np.sqrt(2 / 64), np.zeros(64)
    W2, b2 = rng.normal(0, 1, (64, 10)) * np.sqrt(2 / 64), np.zeros(10)
    shuffle_rng, updates = np.random.default_rng(0), 0
    for _ in range(epochs):
        order = shuffle_rng.permutation(len(X_train))
        for start in range(0, len(X_train), batch_size):
            idx = order[start:start + batch_size]
            xb, yb = X_train[idx], y_train[idx]
            Z1 = xb @ W1 + b1
            A1 = np.maximum(0, Z1)
            Z2 = A1 @ W2 + b2
            P = np.exp(Z2 - Z2.max(axis=1, keepdims=True))
            P /= P.sum(axis=1, keepdims=True)
            dZ2 = P.copy()
            dZ2[np.arange(len(yb)), yb] -= 1
            dZ2 /= len(yb)
            dZ1 = (dZ2 @ W2.T) * (Z1 > 0)
            W2 -= lr * A1.T @ dZ2
            b2 -= lr * dZ2.sum(axis=0)
            W1 -= lr * xb.T @ dZ1
            b1 -= lr * dZ1.sum(axis=0)
            updates += 1
    pred = (np.maximum(0, X_val @ W1 + b1) @ W2 + b2).argmax(axis=1)
    return updates, np.mean(pred == y_val)

for bs in [8, 32, 128]:
    updates, acc = run(bs)
    print(f"batch {bs:>3}: {updates:>4} updates, validation accuracy {acc:.3f}")

batch   8: 3380 updates, validation accuracy 0.980
batch  32:  860 updates, validation accuracy 0.969
batch 128:  220 updates, validation accuracy 0.956