PyTorch Early Stopping

Early stopping ends training automatically when the model stops improving on the validation set. This prevents the model from memorizing training data in later epochs — a key cause of overfitting. It also saves compute time by not running epochs that produce no benefit.

The Problem Early Stopping Solves

Epoch  Train Loss   Val Loss
  10     0.42         0.45   ← both improving, keep training
  20     0.21         0.26   ← still improving
  30     0.11         0.27   ← val loss plateaus (1st warning)
  40     0.06         0.31   ← val loss worsens (overfitting)
  50     0.03         0.38   ← even worse
                             ← should have stopped at epoch 25!

Early Stopping Logic

best_val_loss = ∞
patience = 10          ← epochs to wait before stopping
no_improve_count = 0

each epoch:
  compute val_loss
  if val_loss < best_val_loss:
    best_val_loss = val_loss
    save model weights
    no_improve_count = 0
  else:
    no_improve_count += 1

  if no_improve_count >= patience:
    stop training
    load saved weights (best model)

Early Stopping Class Implementation

import torch

class EarlyStopping:
    def __init__(self, patience=10, min_delta=0.0, path='best_model.pth'):
        self.patience      = patience
        self.min_delta     = min_delta   # minimum improvement to count as progress
        self.path          = path
        self.best_loss     = float('inf')
        self.counter       = 0
        self.stop_training = False

    def __call__(self, val_loss, model):
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter   = 0
            torch.save(model.state_dict(), self.path)
            print(f"  → Saved best model (val_loss={val_loss:.4f})")
        else:
            self.counter += 1
            print(f"  → No improvement ({self.counter}/{self.patience})")
            if self.counter >= self.patience:
                self.stop_training = True
                print("Early stopping triggered.")

Using Early Stopping in Training

import torch.nn as nn
import torch.optim as optim

model     = nn.Linear(10, 1)
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn   = nn.MSELoss()
stopper   = EarlyStopping(patience=10, path='checkpoint.pth')

for epoch in range(500):   # max 500 epochs
    # Training
    model.train()
    # ... training loop here ...
    train_loss = 0.0  # placeholder

    # Validation
    model.eval()
    with torch.no_grad():
        val_loss = 0.0  # placeholder for actual val_loss

    print(f"Epoch {epoch+1}  Train: {train_loss:.4f}  Val: {val_loss:.4f}")

    stopper(val_loss, model)
    if stopper.stop_training:
        print(f"Stopped at epoch {epoch+1}")
        break

# Restore best weights
model.load_state_dict(torch.load('checkpoint.pth'))
model.eval()

Choosing Patience

patience = 5:
  Stops quickly; may stop too early during a temporary plateau
  Use when training is very fast (small dataset)

patience = 10:
  Standard choice; balances responsiveness and stability
  Good for most tasks

patience = 20–30:
  More forgiving; allows the model to escape local plateaus
  Use when training is slow (large dataset, deep model)

min_delta — Minimum Improvement

The min_delta parameter sets a minimum improvement threshold. A decrease in validation loss smaller than min_delta does not count as improvement. This prevents stopping from being triggered by tiny fluctuations that look like progress but aren't meaningful.

stopper = EarlyStopping(patience=10, min_delta=0.001)
# Only counts as improvement if val_loss drops by at least 0.001

Monitoring Accuracy Instead of Loss

For classification tasks, you may prefer to monitor validation accuracy. Simply change the logic to save when accuracy goes up instead of when loss goes down.

class AccuracyEarlyStopping:
    def __init__(self, patience=10, path='best.pth'):
        self.patience     = patience
        self.best_acc     = 0.0
        self.counter      = 0
        self.path         = path
        self.stop         = False

    def __call__(self, val_acc, model):
        if val_acc > self.best_acc:
            self.best_acc = val_acc
            self.counter  = 0
            torch.save(model.state_dict(), self.path)
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.stop = True

Early Stopping with a Learning Rate Scheduler

Combine early stopping with ReduceLROnPlateau for a complete training strategy. The scheduler reduces the learning rate when improvement stalls. Early stopping ends training if the model still doesn't improve after the rate reduction.

from torch.optim.lr_scheduler import ReduceLROnPlateau

scheduler = ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
stopper   = EarlyStopping(patience=15)

for epoch in range(500):
    # ... train ...
    val_loss = evaluate(model, val_loader)
    scheduler.step(val_loss)   # reduce lr if plateau
    stopper(val_loss, model)   # stop if no improvement
    if stopper.stop_training:
        break

Summary

Early stopping monitors validation loss (or accuracy) after each epoch and stops training when performance stops improving. It requires a patience hyperparameter — the number of epochs to wait before stopping. When a new best validation score is reached, save the model weights. When training stops, restore the best saved weights. Combine early stopping with ReduceLROnPlateau for a robust training strategy that avoids overfitting without manually tuning the number of epochs. Early stopping is one of the simplest and most effective tools in the deep learning practitioner's toolkit.

Leave a Comment

Your email address will not be published. Required fields are marked *