PyTorch Train, Validation and Test Splits
A three-way split divides data into train, validation, and test slices. Each slice has one job, and the test slice is used once, at the very end.
One shuffle, three slices
randperm shuffles the row numbers once. Consecutive parts of that one order become the three slices, so no row lands twice.
torch.manual_seed(0)
order = torch.randperm(10)
print(order[:6])
print(order[6:8])
print(order[8:])Why three, not two
Two slices force one slice to do two jobs. Settings get tuned on the same rows that give the final score, and that score flatters the model.
jobs = ["train learns", "validation guides", "test judges once"]
for job in jobs:
print(job)Example
import torch
torch.manual_seed(0)
rows = 300
x = torch.randn(rows, 4)
y = x.sum(dim=1)
order = torch.randperm(rows)
train_idx = order[:180]
val_idx = order[180:240]
test_idx = order[240:]
train_x, train_y = x[train_idx], y[train_idx]
val_x, val_y = x[val_idx], y[val_idx]
test_x, test_y = x[test_idx], y[test_idx]
print("Train rows:", len(train_x))
print("Validation rows:", len(val_x))
print("Test rows:", len(test_x))
print("Total rows:", len(train_x) + len(val_x) + len(test_x))
overlap = len(set(train_idx.tolist()) & set(test_idx.tolist()))
print("Rows shared by train and test:", overlap)
print("First test target:", f"{test_y[0].item():.4f}")Key points
- Train slice: the rows the model learns from
- Validation slice: the rows used to compare settings and watch progress
- Test slice: the rows kept untouched until one final check
- Use torch.randperm to shuffle row numbers once before slicing
- One seed makes the same shuffle, and the same split, every run
In this example, one seed makes the split repeatable. The data is 300 rows of random inputs, and each target is the sum of its row. randperm shuffles the 300 row numbers once. The first 180 numbers form the train slice. The next 60 form the validation slice. The last 60 form the test slice, which the program never trains on. The sizes print first, and they add back up to 300. A set overlap check confirms train and test share zero rows. The first test target prints last, read only at the end.
Note
Tuning settings against the test slice leaks its answers into the choices, so its score stops being honest
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
When should the test slice be used?