How to implement early stopping in Tensorflow, Keras and Pytorch.

In Pytorch we don’t have a built-in early-stopping callback function like Keras. However, we can easily implement early stopping manually in PyTorch. The simple hack is to track the validation loss over the epochs and if it doesn’t improve after a certain number of epochs, we must stop the training

Let’s start coding.

import torch
from torch import nn, optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np

# Sample dataset and model
X_train = np.random.rand(1000, 28 * 28).astype(np.float32)
y_train = np.random.randint(0, 10, 1000).astype(np.int64)
X_val = np.random.rand(200, 28 * 28).astype(np.float32)
y_val = np.random.randint(0, 10, 200).astype(np.int64)

train_data = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))
val_data = TensorDataset(torch.tensor(X_val), torch.tensor(y_val))

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
val_loader = DataLoader(val_data, batch_size=64)

# Define a simple model
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 10)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

#------------------early stopping implementation —---------------------

# Early stopping parameters
patience = 3
best_val_loss = float('inf')
epochs_without_improvement = 0

# Training loop
for epoch in range(50):
    model.train()
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    
    # Validation phase
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            val_loss += loss.item()
    
    val_loss /= len(val_loader)
    print(f'Epoch {epoch+1}, Validation Loss: {val_loss}')

    # Early stopping condition
    if val_loss < best_val_loss: best_val_loss = val_loss epochs_without_improvement = 0 # Save the model checkpoint torch.save(model.state_dict(), 'best_model.pth') else: epochs_without_improvement += 1 if epochs_without_improvement >= patience:
        print("Early stopping...")
        break

Output

Epoch 1, Validation Loss: 2.345960259437561
Epoch 2, Validation Loss: 2.3512794971466064
Epoch 3, Validation Loss: 2.3335058093070984
Epoch 4, Validation Loss: 2.3214242458343506
Epoch 5, Validation Loss: 2.3485482335090637
Epoch 6, Validation Loss: 2.3570011854171753
Epoch 7, Validation Loss: 2.3309141993522644
Early stopping...

As we can see the model training was stopped as val_loss increased

Key Points:

• Track val_loss or another metric for validation performance.
• Stop training when the metric stops improving for a certain number of epochs (patience).
• Save the best model weights during training and load them after stopping.