PyTorch K-Fold Cross-Validation

K-fold cross-validation rotates the validation slice through the data k times and averages the scores. Every row validates exactly once, so small datasets give a steadier score.

The rotation

Thirty rows cut five ways give folds of six rows. Each fold validates once while the other 24 rows train.

rows = 30
k = 5
fold_size = rows // k
for fold in range(k):
    start = fold * fold_size
    print(f"Fold {fold + 1} validates rows {start} to {start + fold_size - 1}")

Average the fold scores

Five scores become one by the plain mean. One lucky split can no longer carry the verdict alone.

scores = [0.967, 1.000, 1.000, 0.967, 0.967]
print(f"{sum(scores) / len(scores):.3f}")

Example

import torch
from torch import nn

torch.manual_seed(0)
per_class = 50
centres = [(-2.0, 0.0), (2.0, 0.0), (0.0, 2.5)]
parts_x, parts_y = [], []
for label, (cx, cy) in enumerate(centres):
    parts_x.append(torch.randn(per_class, 2) * 0.6 + torch.tensor([cx, cy]))
    parts_y.append(torch.full((per_class,), label, dtype=torch.long))
x = torch.cat(parts_x)
y = torch.cat(parts_y)
order = torch.randperm(len(x))
x, y = x[order], y[order]

k = 5
fold_size = len(x) // k
scores = []
for fold in range(k):
    val_idx = list(range(fold * fold_size, (fold + 1) * fold_size))
    train_idx = [i for i in range(len(x)) if i not in val_idx]
    train_x, train_y = x[train_idx], y[train_idx]
    val_x, val_y = x[val_idx], y[val_idx]
    torch.manual_seed(1)
    model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 3))
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.Adam(model.parameters(), lr=0.05)
    for epoch in range(40):
        loss = loss_fn(model(train_x), train_y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    with torch.no_grad():
        pred = model(val_x).argmax(dim=1)
    acc = (pred == val_y).float().mean().item()
    scores.append(acc)
    print(f"Fold {fold + 1}: accuracy {acc:.3f} on {len(val_x)} rows")

mean = sum(scores) / len(scores)
print(f"Mean accuracy: {mean:.3f}")

Key points

Fold: one of the k equal parts the data is cut into
Each fold takes a turn as the validation slice
The other folds train the model for that round
Use a fresh model for every fold so no round inherits trained weights
The mean of the fold scores is the reported score

In this example, one seed builds 150 rows in three blobs. Each blob holds 50 rows around its own centre. A shuffle mixes the classes before the folds are cut. Five folds of 30 rows follow from the shuffled order. Each round trains a fresh model on 120 rows for 40 epochs. The seed resets before each model, so every round starts from the same kind of start. The held-out fold then scores the model. The fold scores print one per line, from 0.967 to 1.000. Their mean prints last at 0.980.

Note

K-fold costs k training runs, so it suits small data better than large data

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn

torch.manual_seed(0)
per_class = 50
centres = [(-2.0, 0.0), (2.0, 0.0), (0.0, 2.5)]
parts_x, parts_y = [], []
for label, (cx, cy) in enumerate(centres):
    parts_x.append(torch.randn(per_class, 2) * 0.6 + torch.tensor([cx, cy]))
    parts_y.append(torch.full((per_class,), label, dtype=torch.long))
x = torch.cat(parts_x)
y = torch.cat(parts_y)
order = torch.randperm(len(x))
x, y = x[order], y[order]

k = 5
fold_size = len(x) // k
scores = []
for fold in range(k):
    val_idx = list(range(fold * fold_size, (fold + 1) * fold_size))
    train_idx = [i for i in range(len(x)) if i not in val_idx]
    train_x, train_y = x[train_idx], y[train_idx]
    val_x, val_y = x[val_idx], y[val_idx]
    torch.manual_seed(1)
    model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 3))
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.Adam(model.parameters(), lr=0.05)
    for epoch in range(40):
        loss = loss_fn(model(train_x), train_y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    with torch.no_grad():
        pred = model(val_x).argmax(dim=1)
    acc = (pred == val_y).float().mean().item()
    scores.append(acc)
    print(f"Fold {fold + 1}: accuracy {acc:.3f} on {len(val_x)} rows")

mean = sum(scores) / len(scores)
print(f"Mean accuracy: {mean:.3f}")

Exercise

In 5-fold cross-validation, how many times does each row serve as validation?

Five times Never Exactly once Once per epoch