🎓 A1Learner
🔍

PyTorch The Eval Loop Pattern

The eval loop pattern measures a model on batched data without training it. It weights each batch by its row count, so a short last batch cannot distort the average.

The trap in plain numbers

Two batch means look easy to average. The plain mean treats a batch of four and a batch of two as equals, and the answer shifts.

batch_means = [0.90, 0.30]
batch_sizes = [4, 2]
naive = sum(batch_means) / len(batch_means)
weighted = sum(m * n for m, n in zip(batch_means, batch_sizes)) / sum(batch_sizes)
print(f"Naive: {naive:.2f}")
print(f"Weighted: {weighted:.2f}")

One loss per row

The argument reduction="none" stops the loss function from averaging. Two rows in give a tensor of two losses out.

torch.manual_seed(1)
model = nn.Linear(3, 2)
loss_fn = nn.CrossEntropyLoss(reduction="none")
x = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
y = torch.tensor([0, 1])
model.eval()
with torch.no_grad():
    losses = loss_fn(model(x), y)
print(losses.shape)

Example

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

torch.manual_seed(0)
x = torch.randn(10, 3)
y = (x[:, 0] + x[:, 1] > 0).long()

torch.manual_seed(1)
model = nn.Linear(3, 2)
loss_fn = nn.CrossEntropyLoss(reduction="none")

data = TensorDataset(x, y)
loader = DataLoader(data, batch_size=4)

model.eval()
total_loss = 0.0
total_rows = 0
batch_means = []
with torch.no_grad():
    for batch_x, batch_y in loader:
        losses = loss_fn(model(batch_x), batch_y)
        batch_means.append(losses.mean().item())
        total_loss += losses.sum().item()
        total_rows += len(batch_y)
        print(f"Batch of {len(batch_y)} rows: mean loss {losses.mean().item():.4f}")

weighted = total_loss / total_rows
naive = sum(batch_means) / len(batch_means)
print(f"Weighted average loss: {weighted:.4f}")
print(f"Naive average loss: {naive:.4f}")
print("Averages differ:", abs(weighted - naive) > 0.001)

Key points

  • Use model.eval() to switch off training behaviour before measuring
  • Use torch.no_grad() to skip gradient tracking during evaluation
  • Weighted average: each batch counts in proportion to its rows
  • Use reduction="none" to get one loss per row from a loss function
  • Sum the per-row losses, then divide by the total row count

In this example, ten rows are scored in batches of four. The batches hold four, four, and two rows. The loss function uses reduction="none", so each row keeps its own loss. The loop adds every row loss into one running total. It also stores each batch mean for the naive average. Each batch prints its size and its mean loss as it passes. The weighted average divides the total loss by all ten rows and gives 0.8130. The naive average of the three batch means gives 0.7529. The short last batch pulled the naive number down, and the final line confirms the two averages differ.

Note

CrossEntropyLoss averages rows inside one call, so the trap appears only when averaging batch means by hand

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does the eval loop weight batches by row count?