Minibatches, Datasets & DataLoaders

So far we've pushed the whole training set through the network in one go. That works for a thousand points; it collapses for a million images — they simply don't fit in memory. The fix is to train on small minibatches, and PyTorch has a dedicated data pipeline for it: Dataset and DataLoader. By the end of this lesson you'll have a complete, reusable training recipe — validation per epoch, best-model checkpointing, and early stopping — that carries you through the rest of the course.

Why minibatches?

Instead of computing the loss on all N examples per update, split the data into chunks of, say, 64 and update after each chunk. Three wins:

1. Memory. Only one batch lives on the GPU at a time. A dataset of any size trains on a fixed memory budget.
2. Faster convergence. With batches of 64 and 64,000 examples, you get 1,000 weight updates per epoch instead of one. Each update uses a noisier gradient estimate, but a thousand decent steps beat one perfect step.
3. Noise as a feature. Batch gradients only approximate the full gradient, so the path downhill jitters. That jitter acts as a mild regularizer and helps the optimizer escape sharp, poorly-generalizing minima.

Three words you'll now use precisely:

• Batch — one chunk of examples (its size is the batch size).
• Iteration — one forward-backward-step cycle on one batch.
• Epoch — one full pass over the dataset. With 1,000 samples and batch size 64, an epoch is ceil(1000 / 64) = 16 iterations.

Dataset and DataLoader

PyTorch splits the job in two. A Dataset answers two questions: how many examples are there? (__len__) and give me example i (__getitem__). A DataLoader wraps a dataset and handles batching, shuffling, and parallel loading. Writing a custom dataset is just the OOP you learned last lesson:

import torch
from torch.utils.data import Dataset, DataLoader, TensorDataset
 
class MyDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)
 
    def __len__(self):
        return len(self.X)              # how many samples?
 
    def __getitem__(self, i):
        return self.X[i], self.y[i]     # one (features, label) pair
 
# For tensors that already exist, TensorDataset does the same thing:
import numpy as np
rng = np.random.default_rng(42)
X, y = rng.normal(size=(1000, 8)), rng.integers(0, 2, 1000)
 
train_set = TensorDataset(
    torch.tensor(X, dtype=torch.float32),
    torch.tensor(y, dtype=torch.long),
)
 
trainloader = DataLoader(train_set, batch_size=64, shuffle=True)
 
for xb, yb in trainloader:              # 16 iterations per epoch
    print(xb.shape, yb.shape)           # torch.Size([64, 8]) torch.Size([64])
    break

Custom Dataset classes earn their keep when __getitem__ does real work — loading an image file from disk, applying transforms — so that only the current batch is ever in memory.

Two DataLoader knobs matter most:

• shuffle=True for training, always. If the data is sorted (all class 0, then all class 1...), each batch is lopsided and the gradients lurch from one class's preferences to the other's. Shuffling each epoch keeps batches representative. For validation, shuffle=False — order doesn't affect a metric, and reproducibility is nice.
• batch_size trades speed for noise. Small batches (8–32): noisy gradients, more regularization, slower wall-clock. Large batches (256+): smooth gradients, better hardware utilization, but more memory and sometimes worse generalization. 64 or 128 is a sensible default; cut it if you hit out-of-memory errors.

The full modern recipe

Minibatches change the loop's shape: an inner loop over batches, nested in an outer loop over epochs — and after each epoch, an evaluation pass over a validation set the model never trains on. Comparing the two curves is how you diagnose overfitting, and the validation score drives two more upgrades:

• Checkpointing — whenever validation improves, save the weights with torch.save(model.state_dict(), path). Training can wander into overfitting; the best model is safely on disk.
• Early stopping — if validation hasn't improved for patience epochs in a row, stop. No babysitting, no wasted compute, no guessing the right number of epochs in advance.

Here's the whole thing as a reusable function — copy it into your projects:

import torch
from torch import nn, optim
 
device = "cuda" if torch.cuda.is_available() else "cpu"
 
def run_epoch(model, loader, criterion, optimizer=None):
    """One pass over loader. Trains if optimizer is given, else evaluates."""
    training = optimizer is not None
    model.train() if training else model.eval()
 
    total_loss, correct, n = 0.0, 0, 0
    with torch.enable_grad() if training else torch.inference_mode():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            output = model(xb)
            loss = criterion(output, yb)
 
            if training:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
 
            total_loss += loss.item() * len(xb)      # weight by batch size
            correct += (output.argmax(1) == yb).sum().item()
            n += len(xb)
    return total_loss / n, correct / n
 
def train(model, trainloader, valloader, criterion, optimizer,
          max_epochs=200, patience=10, path="best_model.pth"):
    best_val_loss, wait = float("inf"), 0
 
    for epoch in range(1, max_epochs + 1):
        train_loss, train_acc = run_epoch(model, trainloader, criterion, optimizer)
        val_loss, val_acc = run_epoch(model, valloader, criterion)
 
        if val_loss < best_val_loss:                 # improvement: checkpoint
            best_val_loss, wait = val_loss, 0
            torch.save(model.state_dict(), path)
        else:                                        # no improvement
            wait += 1
            if wait >= patience:
                print(f"early stopping at epoch {epoch}")
                break
 
        print(f"epoch {epoch:3d} | train loss {train_loss:.4f} acc {train_acc:.3f}"
              f" | val loss {val_loss:.4f} acc {val_acc:.3f}")
 
    model.load_state_dict(torch.load(path))          # restore the best weights
    return model

Details worth noticing:

• The last batch is usually smaller than batch_size, so we accumulate loss.item() * len(xb) and divide by the total count — a plain average of batch losses would weight the final stragglers too heavily.
• run_epoch serves both phases: pass an optimizer to train, omit it to evaluate under inference_mode.
• We monitor validation loss and reload the checkpoint at the end, so the function returns the best model, not the last one.

And using it is three lines:

model = nn.Sequential(
    nn.Linear(8, 32), nn.ReLU(), nn.Dropout(0.2),
    nn.Linear(32, 2),
).to(device)
 
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
 
model = train(model, trainloader, valloader, criterion, optimizer)

Next up: a new module — convolutional neural networks, where we finally stop flattening images and let the network see their 2-D structure.