04 · Splits & Built-in Datasets ✂️

5 min read

Before you train, you must divide your data honestly — and for learning and prototyping, PyTorch hands you famous datasets ready to go.

1. The 'Why'

A model that scores 100% on the data it was trained on has told you almost nothing — it may have simply memorized the answers. The only way to know whether a model has actually learned something that generalizes is to evaluate it on data it has never seen during training. That's why we split a dataset into a training set (used to update the weights), a validation set (used to tune choices like learning rate and to watch for overfitting), and a test set (touched only once, at the very end, to report honest final performance). Getting this split right is not a formality — it's the foundation of trustworthy results, and it's the piece that makes the training loop in Module 05 meaningful rather than self-congratulatory.

The second half of this sub-module is a gift: you don't always have to build a Dataset from scratch. torchvision, torchtext, and torchaudio ship built-in datasets — MNIST, CIFAR-10, FashionMNIST, and many more — that download themselves and already implement the Dataset interface. They're perfect for learning, prototyping an architecture, and sanity-checking your training loop before pointing it at your own messy data. Combined with random_split for carving out a validation set, these tools let you go from zero to a fully-batched, properly-split data pipeline in a handful of lines — exactly the pipeline you'll feed into Module 05.

2. Core Concepts

random_split. torch.utils.data.random_split(dataset, lengths) randomly partitions a dataset into non-overlapping subsets. Pass a generator with a fixed seed for reproducible splits.

Why three splits. Train fits the weights; validation guides your decisions during development (and detects overfitting); test is a one-time final exam. Never tune on the test set — that leaks information and inflates your reported score.

Built-in datasets. Classes like torchvision.datasets.MNIST accept a root folder, train=True/False to pick the official split, download=True to fetch it, and a transform. They return the same (image, label) samples your custom Dataset would.

One transform per split. Apply your augmenting train_tfms to the training subset and the deterministic val_tfms to validation/test (sub-module 03). With random_split you typically set the transform on the underlying dataset(s) accordingly.

3. Code in Action

Splitting a dataset into train / validation

import torch
from torch.utils.data import random_split, DataLoader

# Suppose `full_ds` is any Dataset of 1000 samples
full_ds = ...                              # your custom or built-in dataset
n_val = int(0.2 * len(full_ds))            # hold out 20% for validation
n_train = len(full_ds) - n_val

gen = torch.Generator().manual_seed(42)    # fixed seed -> reproducible split
train_ds, val_ds = random_split(full_ds, [n_train, n_val], generator=gen)

print(len(train_ds), len(val_ds))          # 800 200

Loading a built-in dataset (MNIST) end-to-end

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

tfms = transforms.Compose([
    transforms.ToTensor(),                             # PIL -> tensor in [0,1]
    transforms.Normalize((0.1307,), (0.3081,)),        # MNIST's known mean/std
])

# Official train and test splits download themselves on first run
train_ds = datasets.MNIST(root="./data", train=True,  download=True, transform=tfms)
test_ds  = datasets.MNIST(root="./data", train=False, download=True, transform=tfms)

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_ds,  batch_size=64, shuffle=False)

images, labels = next(iter(train_loader))
print(images.shape, labels.shape)          # torch.Size([64, 1, 28, 28]) torch.Size([64])

The complete pipeline, ready for Module 05

# This is the exact handoff into the training loop:
#   for epoch in range(EPOCHS):
#       for batch_X, batch_y in train_loader:   # <-- built here
#           ...                                  # forward, loss, backward, step
#       # then evaluate on val_loader / test_loader
print("Train batches:", len(train_loader), "| Test batches:", len(test_loader))

4. Common Pitfalls

Pitfall 1 — Tuning on the test set. If you repeatedly check test accuracy and adjust your model to improve it, the test set stops being an honest measure — you've effectively trained on it. Use the validation set for all decision-making and reserve the test set for a single final evaluation.

Pitfall 2 — Non-reproducible splits. Calling random_split without a seeded generator gives a different split every run, making results impossible to compare. Always pass generator=torch.Generator().manual_seed(...).

Pitfall 3 — Data leakage across splits. If you normalize using statistics computed over the whole dataset, or let augmented copies of the same image land in both train and test, information leaks and your metrics lie. Compute normalization stats from the training set only, and split before augmenting.

5. Further Reading & Watch List

⬅️ Prev: 03 · Transforms · 🏁 Module complete! Back to the module index or on to Module 05 — The Training Loop.