🎓 A1Learner
🔍

PyTorch Overfitting, a First Look

Overfitting is learning the training rows by heart instead of the pattern. The training loss keeps falling, and the validation loss starts to climb.

A big model, a tiny dataset

The model holds far more parameters than the dataset holds rows. Memorising the rows is within its reach.

count = sum(p.numel() for p in model.parameters())
print("Parameters:", count)
print("Training rows:", len(x_train))

Both losses before training

Two scores start the story. The training rows carry noise, and the validation rows follow the clean curve.

model.eval()
with torch.no_grad():
    train_loss = loss_fn(model(x_train), y_train)
    val_loss = loss_fn(model(x_val), y_val)
print(f"Train: {train_loss.item():.4f}")
print(f"Val: {val_loss.item():.4f}")

Noise in the training rows

The training targets add noise to the true curve. A model that fits the noise fits rows, not the rule.

print("True curve at x=0:", 0.0 ** 2)
print("Noise can move a target up or down")

Example

import torch
import torch.nn as nn

torch.manual_seed(4)
g = torch.Generator().manual_seed(4)
x_train = torch.linspace(-1, 1, 10).reshape(-1, 1)
y_train = x_train ** 2 + torch.randn(10, 1, generator=g) * 0.3
x_val = torch.linspace(-1, 1, 21).reshape(-1, 1)
y_val = x_val ** 2

model = nn.Sequential(
    nn.Linear(1, 32),
    nn.ReLU(),
    nn.Linear(32, 32),
    nn.ReLU(),
    nn.Linear(32, 1),
)
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

print("Epoch  Train   Val")
for epoch in range(201):
    if epoch % 50 == 0:
        model.eval()
        with torch.no_grad():
            train_loss = loss_fn(model(x_train), y_train)
            val_loss = loss_fn(model(x_val), y_val)
        print(f"{epoch:3d}    {train_loss.item():.4f}  {val_loss.item():.4f}")
    if epoch == 200:
        break
    model.train()
    optimizer.zero_grad()
    loss = loss_fn(model(x_train), y_train)
    loss.backward()
    optimizer.step()

Key points

  • Overfitting: the model learns the training rows by heart and does worse on new rows
  • Training loss falls while validation loss climbs in an overfitting run
  • A model with far more parameters than rows can memorise the rows
  • Validation loss is the signal that separates learning from memorising
  • A low training loss alone does not prove the model learned the pattern

In this example, 10 noisy training rows follow the curve y = x squared. The validation rows follow the same curve with no noise. The model is wide, with two hidden layers of 32 units. Its training loss falls at every printed epoch. The validation loss falls until epoch 50 and then climbs. The model keeps improving on the rows it sees and gets worse on the clean rows, which is overfitting.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A run shows training loss falling and validation loss climbing. What is happening?