Early stopping is a regularization technique that halts model training when validation performance ceases to improve, thereby preventing overfitting. This approach monitors validation loss across epochs and terminates training if no significant improvement occurs for a predetermined number of epochs.

Early Stopping Implementation

The following class provides a reusable early stopping mechanism:

class TrainingHaltChecker:
    def __init__(self, wait_epochs=5, threshold=0.001):
        """
        Parameters:
        - wait_epochs: Number of epochs to wait before stopping (patience)
        - threshold: Minimum improvement required to reset the wait counter
        """
        self.wait_epochs = wait_epochs
        self.threshold = threshold
        self.stagnant_epochs = 0
        self.optimal_loss = None
        self.should_halt = False

    def __call__(self, validation_loss):
        if self.optimal_loss is None:
            self.optimal_loss = validation_loss  # Initialize on first call
        elif validation_loss < self.optimal_loss - self.threshold:
            self.optimal_loss = validation_loss
            self.stagnant_epochs = 0  # Reset counter on improvement
        else:
            self.stagnant_epochs += 1
            if self.stagnant_epochs >= self.wait_epochs:
                self.should_halt = True  # Trigger stopping condition

Complete Training Example with Fashion-MNIST

This example demonstrates integrating early stopping into a convolutional neural network training pipeline:

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
from tqdm import tqdm

# Define a simple CNN architecture
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    
    def forward(self, x):
        return self.fc_layers(self.conv_layers(x))

# Load and prepare Fashion-MNIST dataset
train_dataset = datasets.FashionMNIST(
    root="./fashion_data",
    train=True,
    download=True,
    transform=ToTensor()
)

val_dataset = datasets.FashionMNIST(
    root="./fashion_data",
    train=False,
    download=True,
    transform=ToTensor()
)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)

# Initialize model components
device = "cuda" if torch.cuda.is_available() else "cpu"
network = ImageClassifier().to(device)
criterion = nn.CrossEntropyLoss()
solver = torch.optim.Adam(network.parameters(), lr=0.001)

# Training epoch function
def training_epoch(loader, net, loss_func, optimizer):
    net.train()
    total_samples = len(loader.dataset)
    batch_count = len(loader)
    cumulative_loss = 0.0
    correct_predictions = 0
    
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        
        outputs = net(images)
        loss = loss_func(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        cumulative_loss += loss.item()
        correct_predictions += (outputs.argmax(dim=1) == labels).float().sum().item()
    
    avg_loss = cumulative_loss / batch_count
    accuracy = correct_predictions / total_samples
    return avg_loss, accuracy

# Validation epoch function
def validation_epoch(loader, net, loss_func):
    net.eval()
    total_samples = len(loader.dataset)
    batch_count = len(loader)
    cumulative_loss = 0.0
    correct_predictions = 0
    
    with torch.no_grad():
        for images, labels in tqdm(loader, desc="Validating"):
            images, labels = images.to(device), labels.to(device)
            
            outputs = net(images)
            loss = loss_func(outputs, labels)
            
            cumulative_loss += loss.item()
            correct_predictions += (outputs.argmax(dim=1) == labels).float().sum().item()
    
    avg_loss = cumulative_loss / batch_count
    accuracy = correct_predictions / total_samples
    return avg_loss, accuracy

# Training loop with early stopping
halt_checker = TrainingHaltChecker(wait_epochs=3)
max_epochs = 50
best_val_acc = 0.0

history = {
    'train_loss': [],
    'train_acc': [],
    'val_loss': [],
    'val_acc': []
}

for epoch in range(max_epochs):
    print(f"\nEpoch {epoch+1}/{max_epochs}")
    
    tr_loss, tr_acc = training_epoch(train_loader, network, criterion, solver)
    val_loss, val_acc = validation_epoch(val_loader, network, criterion)
    
    history['train_loss'].append(tr_loss)
    history['train_acc'].append(tr_acc)
    history['val_loss'].append(val_loss)
    history['val_acc'].append(val_acc)
    
    # Save best performing model
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(network.state_dict(), "best_classifier.pth")
        print("→ Saved new best model")
    
    # Check early stopping condition
    halt_checker(val_loss)
    if halt_checker.should_halt:
        print("→ Early stopping triggered: validation loss plateaued")
        break
    
    print(f"Train Loss: {tr_loss:.4f} | Train Acc: {tr_acc:.4f}")
    print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}")

# Plot training curves
plt.style.use('seaborn-v0_8-darkgrid')
fig, axes = plt.subplots(1, 2, figsize=(12, 4))

# Loss subplot
axes[0].plot(history['train_loss'], label='Training Loss', linewidth=2)
axes[0].plot(history['val_loss'], label='Validation Loss', linewidth=2)
axes[0].set_xlabel('Epoch', fontsize=11)
axes[0].set_ylabel('Loss', fontsize=11)
axes[0].set_title('Loss Evolution', fontsize=12, fontweight='bold')
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# Accuracy subplot
axes[1].plot(history['train_acc'], label='Training Accuracy', linewidth=2)
axes[1].plot(history['val_acc'], label='Validation Accuracy', linewidth=2)
axes[1].set_xlabel('Epoch', fontsize=11)
axes[1].set_ylabel('Accuracy', fontsize=11)
axes[1].set_title('Accuracy Evolution', fontsize=12, fontweight='bold')
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

This implementation tracks both training and validation metrics, saves the best model based on validation accuracy, and automatically terminates training when validation loss stops improving for three consecutive epochs. The visualization provides clear insights into model performance and the effectiveness of early stopping.