🎓 A1Learner
🔍

PyTorch Early Stopping

Early stopping ends training when the validation loss stops improving. It keeps the weights from the best epoch instead of the last one.

Find the best epoch

The validation loss falls, then rises again. The minimum marks the epoch worth keeping.

val_losses = [0.50, 0.32, 0.25, 0.26, 0.31]
best = min(val_losses)
print("Best loss:", best)
print("Best epoch:", val_losses.index(best) + 1)

Count the wait

Each worse epoch adds one to a counter. The counter reaching patience is the stop signal.

patience = 2
wait = 0
for val in [0.25, 0.26, 0.31]:
    wait += 1
print("Wait after three worse epochs:", wait)
print("Time to stop:", wait >= patience)

Example

import copy
import torch
from torch import 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)

best_val = float("inf")
best_state = None
best_epoch = 0
wait = 0
patience = 25
stopped = 0
final_val = None
for epoch in range(1, 301):
    model.train()
    optimizer.zero_grad()
    loss = loss_fn(model(x_train), y_train)
    loss.backward()
    optimizer.step()
    model.eval()
    with torch.no_grad():
        val = loss_fn(model(x_val), y_val).item()
    final_val = val
    if val < best_val - 1e-4:
        best_val = val
        best_epoch = epoch
        best_state = copy.deepcopy(model.state_dict())
        wait = 0
    else:
        wait += 1
        if wait >= patience:
            stopped = epoch
            break

model.load_state_dict(best_state)
model.eval()
with torch.no_grad():
    restored_val = loss_fn(model(x_val), y_val).item()

print(f"Stopped at epoch: {stopped}")
print(f"Best epoch: {best_epoch}")
print(f"Best val loss: {best_val:.4f}")
print(f"Final val loss: {final_val:.4f}")
print(f"Restored val loss: {restored_val:.4f}")
print("Restored beats final:", restored_val < final_val)

Key points

  • Patience: how many epochs without improvement the run allows before stopping
  • Best state: a deep copy of the state_dict from the best epoch
  • Use copy.deepcopy on the state_dict, because a plain reference changes as training continues
  • Overfitting: the training loss keeps falling while the validation loss rises
  • Restoring the best state gives back the model from the best epoch

In this example, a model with two hidden layers learns a curve from only ten noisy points. The validation rows follow the clean curve, so overfitting shows up fast. Each epoch trains once, then measures the validation loss. A better loss saves a deep copy of the state_dict and resets the wait counter. A worse loss grows the counter, and patience is 25 epochs. The run stops at epoch 66, well before the 300 allowed. The best epoch was 41, with a validation loss of 0.0259. The final weights score 0.0466, and the restored best weights score 0.0259 again. The last line confirms the restored weights beat the final ones.

Note

A small tolerance on improvement stops noise in the loss from resetting the patience counter

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why must the best state_dict be deep-copied when it is saved?