Train, Validation, and Test Splits

Train, validation, and test splits are the guardrails that keep a deep learning experiment honest. The training split is the only data used to update model weights. The validation split is used during development to choose architectures, tune hyperparameters, set early stopping, and compare candidate runs. The test split is held back until the end so it can estimate how the chosen model behaves on examples it did not influence.

In a PyTorch data pipeline, a split is usually represented as a set of integer indices into a Dataset, a wrapped Subset, or a Sampler used by a DataLoader. The outcome you want is not just three filenames or three percentages. You want three non-overlapping views of the same prediction problem, with labels and important groups distributed deliberately, random choices reproducible, and no preprocessing step learning from validation or test data.

How Splits Work Inside PyTorch

A PyTorch Dataset defines two essential methods: __len__, which reports how many examples are addressable, and __getitem__, which returns one example by integer position. Splitting normally does not copy tensors or image files. Instead, PyTorch stores index lists. torch.utils.data.Subset(dataset, indices) keeps a reference to the original dataset and translates local positions back into the original indices. If a subset receives local index 0, it returns dataset[indices[0]].

random_split automates this index partitioning. It creates a permutation of dataset positions and slices that permutation into requested lengths. Passing a torch.Generator with a fixed seed makes the permutation reproducible, which is essential when comparing model changes. Without a fixed generator, two runs may train and validate on different examples, making metric changes hard to interpret.

DataLoader then controls iteration. For training, shuffle=True reshuffles examples each epoch so minibatches vary. For validation and test, shuffle=False keeps evaluation stable and easier to debug. Shuffling the validation loader rarely changes aggregate metrics, but it makes per-example investigation and batch-specific failure reproduction harder.

API Anatomy

The simplest PyTorch API is random_split(dataset, lengths, generator=None). lengths can be explicit counts such as [800, 100, 100]. Recent PyTorch releases also accept fractions in many environments, but explicit counts are clearer in teaching code because they expose rounding decisions. If the dataset has 1,003 examples and you ask for 80/10/10, one split must receive the remainder. Decide that intentionally.

For classification, a purely random split can produce weak validation signals when classes are imbalanced. A validation set with too few rare-class examples may make a model look better or worse by chance. PyTorch core does not provide a stratified splitter, but you can build one by grouping indices by label, shuffling within each label, and assigning a proportional number of examples to each split. For medical, user, document, or time-series data, the key unit may not be an individual row. You may need group-aware or chronological splits so the same patient, customer, document, or future timestamp cannot leak across boundaries.

Example 1: Reproducible Random Split

This first example creates a small tensor dataset, splits it by index, and prints the selected original indices. The important behavior is that the seed controls the split while Subset still points at the original dataset.

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

features = torch.arange(20).float().view(10, 2)
labels = torch.arange(10)
dataset = TensorDataset(features, labels)

train_ds, val_ds, test_ds = random_split(
    dataset,
    [6, 2, 2],
    generator=torch.Generator().manual_seed(7),
)

print("train indices:", train_ds.indices)
print("val indices:", val_ds.indices)
print("test indices:", test_ds.indices)
print("first train example:", train_ds[0][0].tolist(), int(train_ds[0][1]))

The exact indices are deterministic for a given PyTorch permutation implementation and seed. The conceptual output is three disjoint index lists whose lengths are 6, 2, and 2. The first training example is fetched from the original tensors through the selected index, not copied into a new tensor dataset.

Example 2: Stratified Split for Labels

Suppose a binary classification dataset contains eight negative examples and four positive examples. A random split could accidentally put only one positive example in validation and none in test. The next example performs a simple stratified split by label so each split receives examples from both classes.

from collections import defaultdict
import random

labels = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1]
fractions = {"train": 0.5, "val": 0.25, "test": 0.25}

by_label = defaultdict(list)
for index, label in enumerate(labels):
    by_label[label].append(index)

rng = random.Random(3)
splits = {"train": [], "val": [], "test": []}

for label, indices in by_label.items():
    rng.shuffle(indices)
    n = len(indices)
    n_train = round(n * fractions["train"])
    n_val = round(n * fractions["val"])
    splits["train"].extend(indices[:n_train])
    splits["val"].extend(indices[n_train:n_train + n_val])
    splits["test"].extend(indices[n_train + n_val:])

for name in splits:
    splits[name].sort()
    counts = {label: sum(labels[i] == label for i in splits[name]) for label in sorted(set(labels))}
    print(name, splits[name], counts)

The deterministic output contains two classes in each split: training receives four zeros and two ones, validation receives two zeros and one one, and test receives two zeros and one one. This is better for metric stability, but it is still row-level splitting. If rows from the same real-world entity are correlated, stratification alone is not enough.

Example 3: Group-Aware Split

Deep learning datasets often contain multiple records from one entity: several images from one patient, many transactions from one account, or repeated sensor windows from one device. If the same entity appears in both training and test, the model may appear to generalize while merely recognizing entity-specific patterns. A group-aware split assigns whole groups to splits.

from collections import defaultdict

patient_id = ["p1", "p1", "p2", "p2", "p3", "p3", "p4", "p4", "p5", "p5"]
label =      [0,    0,    1,    1,    0,    0,    1,    1,    0,    0]
train_patients = {"p1", "p2", "p3"}
val_patients = {"p4"}
test_patients = {"p5"}

split_indices = defaultdict(list)
for i, pid in enumerate(patient_id):
    if pid in train_patients:
        split_indices["train"].append(i)
    elif pid in val_patients:
        split_indices["val"].append(i)
    elif pid in test_patients:
        split_indices["test"].append(i)
    else:
        raise ValueError(f"unassigned patient: {pid}")

print(dict(split_indices))
print("overlap:", train_patients & val_patients & test_patients)

The output maps all rows for each selected patient into one split. The overlap expression prints an empty set, confirming that no patient appears in all three sets. In real code, check pairwise intersections too: train with validation, train with test, and validation with test.

Example 4: DataLoaders After Splitting

Once splits exist, each one normally receives its own DataLoader. The model sees training minibatches with gradient tracking and optimizer updates. Validation and test run under torch.no_grad() or torch.inference_mode(), with the model in eval() mode so dropout and batch normalization behave as evaluation layers.

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

torch.manual_seed(0)
x = torch.randn(12, 3)
y = (x.sum(dim=1) > 0).long()
dataset = TensorDataset(x, y)
train_ds, val_ds, test_ds = random_split(dataset, [8, 2, 2], generator=torch.Generator().manual_seed(11))

train_loader = DataLoader(train_ds, batch_size=4, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=2, shuffle=False)

for batch_x, batch_y in val_loader:
    print(batch_x.shape, batch_y.shape)

The validation loader prints one batch shaped like torch.Size([2, 3]) torch.Size([2]). The training loader may present examples in a different order each epoch, while the validation loader remains stable.

Design Choices and Trade-Offs

An 80/10/10 split is common, but it is not a law. Small datasets may need cross-validation or repeated splits because one validation slice has high variance. Very large datasets may only need a small percentage for validation and test because absolute example counts are already high. For rare-event classification, choose split sizes that guarantee enough positive examples to make precision, recall, or AUROC meaningful.

Random splits are simple and suitable when examples are independent and identically distributed. Stratified splits protect label proportions. Group-aware splits protect against entity leakage. Chronological splits are usually preferred for forecasting, recommender systems, and logs where the model will be used on future data. The trade-off is that chronological validation may be harder because distributions drift, but that difficulty reflects the real deployment problem.

Preprocessing must respect split boundaries. Fit scalers, tokenizers, imputers, vocabulary builders, and augmentation policies using training data only. Then apply the learned transformation to validation and test. If normalization statistics are computed over the entire dataset before splitting, validation and test information has already influenced training.

Failure Modes and Troubleshooting

Symptom: validation accuracy is extremely high, but real examples fail. Cause: duplicate or near-duplicate examples, users, patients, or time windows appear across splits. Diagnose: hash raw inputs, compare group identifiers, and inspect nearest neighbors between train and test. Correct: split by group or time before augmentation and preprocessing.

Symptom: two training runs with the same code produce incompatible validation curves. Cause: split generation, worker seeding, or loader shuffling is uncontrolled. Diagnose: log split indices, random seeds, and dataset version. Correct: persist split index files or use explicit seeded generators, and compare runs on the same validation indices.

Symptom: validation loss changes when batch size changes, even with the same checkpoint. Cause: the model is still in training mode, so dropout or batch normalization uses training behavior during validation. Diagnose: print model.training before evaluation and run a repeated forward pass on the same batch. Correct: call model.eval() and use torch.inference_mode() for validation and test loops.

Symptom: a class has zero recall on the test set, but the aggregate accuracy looks acceptable. Cause: the split is imbalanced or the metric hides minority-class failure. Diagnose: print per-split label counts and a confusion matrix. Correct: use stratified splitting, collect more minority examples, and report class-aware metrics.

Reliability and Performance Implications

Reliable experiments require split artifacts. Store the dataset version and the exact train, validation, and test indices used for important runs. That lets you rerun a failed experiment, compare architectures fairly, and audit unexpected metric jumps. For large datasets, storing integer index arrays is usually cheaper than copying raw data into three directories.

Performance depends on where splitting happens. Creating Subset objects is cheap because it stores indices. Expensive decoding, resizing, or tokenization still occurs when examples are loaded. If you cache preprocessed data, include the split name or source index in cache keys carefully so validation and test records are not overwritten or mixed with training-only augmentations.

Hands-On Lab

Prerequisites: Python, PyTorch, and a small local dataset or synthetic tensor dataset. You should know the label for each example and, if relevant, a group identifier such as user, patient, or document id.

1. Create a dataset object and print len(dataset), one example shape, one label, and one group id if present.
2. Choose the split policy: random for independent examples, stratified for class imbalance, group-aware for repeated entities, or chronological for future-facing data.
3. Generate train, validation, and test index lists with a fixed seed where randomness is used.
4. Assert that the three index sets are pairwise disjoint and that their combined length equals the dataset length.
5. Print per-split label counts and group counts. Confirm they match the intended evaluation story.
6. Wrap each index list with Subset and create one DataLoader per split.
7. Run one training batch and one validation batch. Verify tensor shapes, dtypes, and that validation runs with model.eval().
8. Save the split indices beside the experiment configuration so future runs can reuse them.

Verification: rerun the script and confirm the saved or printed split indices are unchanged. Confirm no index appears in more than one split. For cleanup, remove temporary split files only after the experiment is no longer needed; otherwise keep them as part of the run record.

Assessment Exercises

1. A dataset has 50,000 images from 2,000 patients. Explain why row-level random splitting can inflate test accuracy and design a better split.
2. You trained five learning rates and picked the checkpoint with the best validation score. Why should the test set not be used again to choose among those five runs?
3. Given class counts of 9,700 negative and 300 positive examples, propose split sizes and metrics that would make validation useful.
4. A normalization transform was fitted before splitting the dataset. Describe the leakage risk and rewrite the pipeline order.
5. Your validation metric changes between runs even though the code and checkpoint are unchanged. List three concrete values or artifacts you would inspect first.

Summary

Train, validation, and test splits are index-level boundaries that shape every PyTorch experiment. The training split teaches the model, validation guides development choices, and test estimates final generalization. Use seeded indices for reproducibility, stratified or group-aware policies when the data demands them, evaluation loaders that do not shuffle, and preprocessing that learns only from training data. A split is correct when it matches the deployment question and can be verified, reproduced, and defended.