Eric Campbell 5 min read Machine Learning

How to Implement Early Stopping in TensorFlow

Early stopping is one of the most effective regularization techniques in deep learning. The core idea is simple: monitor your model's performance on a validation set during training and stop when...

Key Insights

  • Early stopping prevents overfitting by monitoring validation metrics and halting training when performance stops improving, saving computational resources and preventing model degradation
  • TensorFlow’s EarlyStopping callback offers fine-grained control through parameters like patience, restore_best_weights, and min_delta to balance between training thoroughness and efficiency
  • Custom callbacks enable sophisticated stopping criteria beyond single-metric monitoring, allowing you to implement domain-specific logic for optimal model selection

Introduction to Early Stopping

Early stopping is one of the most effective regularization techniques in deep learning. The core idea is simple: monitor your model’s performance on a validation set during training and stop when that performance stops improving. This prevents your model from overfitting to the training data—a scenario where the model memorizes training examples rather than learning generalizable patterns.

Without early stopping, you might train for a fixed number of epochs, potentially wasting computational resources on training that degrades model quality. Worse, you might stop training too soon, before the model has converged. Early stopping solves both problems by dynamically determining the optimal training duration based on actual performance metrics.

The mechanism works by tracking a chosen metric (typically validation loss or accuracy) across epochs. When this metric stops improving for a specified number of epochs (the “patience” period), training halts. Optionally, the model weights can be restored to the point where the metric was optimal, ensuring you get the best-performing model rather than the final one.

Understanding TensorFlow’s EarlyStopping Callback

TensorFlow provides tf.keras.callbacks.EarlyStopping, a built-in callback that integrates seamlessly into the training pipeline. The callback monitors specified metrics after each epoch and applies your stopping criteria automatically.

Key parameters include:

  • monitor: The metric to track (e.g., ‘val_loss’, ‘val_accuracy’)
  • patience: Number of epochs with no improvement before stopping
  • restore_best_weights: Whether to restore model weights from the best epoch
  • min_delta: Minimum change to qualify as an improvement
  • mode: Whether to minimize (‘min’), maximize (‘max’), or auto-detect (‘auto’)
  • baseline: Minimum acceptable value for the monitored metric

Here’s a basic instantiation:

from tensorflow.keras.callbacks import EarlyStopping

early_stopping = EarlyStopping(
    monitor='val_loss',
    patience=5,
    restore_best_weights=True
)

This configuration monitors validation loss, waits 5 epochs for improvement, and restores the best weights when stopping.

Implementing Basic Early Stopping

Let’s implement early stopping in a complete training scenario. We’ll build a simple neural network for binary classification and apply early stopping to prevent overfitting.

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.callbacks import EarlyStopping
import numpy as np

# Generate synthetic data
np.random.seed(42)
X_train = np.random.randn(1000, 20)
y_train = (X_train[:, 0] + X_train[:, 1] > 0).astype(int)
X_val = np.random.randn(200, 20)
y_val = (X_val[:, 0] + X_val[:, 1] > 0).astype(int)

# Define model
model = keras.Sequential([
    layers.Dense(64, activation='relu', input_shape=(20,)),
    layers.Dropout(0.3),
    layers.Dense(32, activation='relu'),
    layers.Dropout(0.3),
    layers.Dense(1, activation='sigmoid')
])

# Compile model
model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy']
)

# Configure early stopping
early_stopping = EarlyStopping(
    monitor='val_loss',
    patience=10,
    restore_best_weights=True,
    verbose=1
)

# Train with early stopping
history = model.fit(
    X_train, y_train,
    validation_data=(X_val, y_val),
    epochs=100,
    batch_size=32,
    callbacks=[early_stopping],
    verbose=1
)

print(f"Training stopped at epoch: {len(history.history['loss'])}")

The verbose=1 parameter in EarlyStopping prints a message when training stops, showing which epoch had the best performance.

Advanced Configuration Options

Different training scenarios require different early stopping configurations. The mode parameter is crucial for metrics where higher is better (like accuracy) versus lower is better (like loss).

# Configuration 1: Stop on validation accuracy plateau
early_stopping_acc = EarlyStopping(
    monitor='val_accuracy',
    mode='max',  # Higher accuracy is better
    patience=15,
    min_delta=0.001,  # Require at least 0.1% improvement
    restore_best_weights=True,
    verbose=1
)

# Configuration 2: Aggressive early stopping for quick experiments
early_stopping_aggressive = EarlyStopping(
    monitor='val_loss',
    mode='min',
    patience=3,  # Stop quickly
    restore_best_weights=False,  # Keep final weights
    verbose=1
)

# Configuration 3: Conservative stopping with baseline
early_stopping_baseline = EarlyStopping(
    monitor='val_loss',
    mode='min',
    patience=20,
    baseline=0.5,  # Only stop if val_loss is below 0.5
    min_delta=0.0001,
    restore_best_weights=True,
    verbose=1
)

The choice between restore_best_weights=True and False depends on your goals. Setting it to True (recommended for most cases) ensures you get the model from the epoch with the best validation performance. Setting it to False gives you the final model, which might be useful if you’re using other regularization techniques that benefit from extended training.

Custom Early Stopping Logic

Sometimes you need more sophisticated stopping criteria than monitoring a single metric. Custom callbacks let you implement complex logic, such as stopping when multiple conditions are met or when metrics exceed dynamic thresholds.

class CustomEarlyStopping(keras.callbacks.Callback):
    def __init__(self, patience=5, min_accuracy=0.85, min_loss=0.3):
        super().__init__()
        self.patience = patience
        self.min_accuracy = min_accuracy
        self.min_loss = min_loss
        self.wait = 0
        self.best_epoch = 0
        self.best_weights = None
        self.best_score = float('inf')
    
    def on_epoch_end(self, epoch, logs=None):
        # Custom metric: weighted combination
        current_score = logs.get('val_loss', 0) - (0.5 * logs.get('val_accuracy', 0))
        val_acc = logs.get('val_accuracy', 0)
        val_loss = logs.get('val_loss', float('inf'))
        
        # Check if both metrics meet minimum thresholds
        if val_acc >= self.min_accuracy and val_loss <= self.min_loss:
            if current_score < self.best_score:
                self.best_score = current_score
                self.best_epoch = epoch
                self.wait = 0
                self.best_weights = self.model.get_weights()
                print(f"\nEpoch {epoch}: Custom metric improved to {current_score:.4f}")
            else:
                self.wait += 1
                if self.wait >= self.patience:
                    self.model.stop_training = True
                    if self.best_weights is not None:
                        self.model.set_weights(self.best_weights)
                    print(f"\nEarly stopping triggered. Restoring weights from epoch {self.best_epoch}")

# Use custom callback
custom_callback = CustomEarlyStopping(patience=7, min_accuracy=0.80, min_loss=0.4)

history = model.fit(
    X_train, y_train,
    validation_data=(X_val, y_val),
    epochs=100,
    callbacks=[custom_callback],
    verbose=0
)

This custom callback only considers stopping once both accuracy and loss meet minimum thresholds, then monitors a weighted combination of both metrics.

Best Practices and Common Pitfalls

Choosing Patience Values: Too low (1-3 epochs) risks stopping before the model converges, especially with noisy validation metrics. Too high (>30 epochs) defeats the purpose. Start with 10-15 for most tasks, adjusting based on your learning rate and dataset size.

Metric Selection: Monitor val_loss for most tasks—it’s more stable than accuracy and directly reflects what you’re optimizing. Use val_accuracy or custom metrics only when they align better with your business objective.

Validation Split Strategy: Ensure your validation set is representative and large enough (typically 10-20% of training data). Small validation sets produce noisy metrics that trigger false stops.

Min Delta Considerations: Set min_delta to filter out noise. For loss, try 0.0001-0.001. For accuracy, use 0.001-0.01 depending on your scale.

Here’s code to visualize where early stopping occurred:

import matplotlib.pyplot as plt

def plot_training_history(history, early_stop_epoch=None):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
    
    # Plot loss
    ax1.plot(history.history['loss'], label='Training Loss')
    ax1.plot(history.history['val_loss'], label='Validation Loss')
    if early_stop_epoch:
        ax1.axvline(x=early_stop_epoch, color='r', linestyle='--', label='Early Stop')
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    ax1.legend()
    ax1.set_title('Model Loss')
    
    # Plot accuracy
    ax2.plot(history.history['accuracy'], label='Training Accuracy')
    ax2.plot(history.history['val_accuracy'], label='Validation Accuracy')
    if early_stop_epoch:
        ax2.axvline(x=early_stop_epoch, color='r', linestyle='--', label='Early Stop')
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Accuracy')
    ax2.legend()
    ax2.set_title('Model Accuracy')
    
    plt.tight_layout()
    plt.savefig('training_history.png', dpi=300, bbox_inches='tight')
    plt.show()

# Use after training
stopped_epoch = len(history.history['loss'])
plot_training_history(history, early_stop_epoch=stopped_epoch)

Combining Early Stopping with Other Callbacks

Early stopping works best when combined with other callbacks for comprehensive training management. ModelCheckpoint saves the best model to disk, ReduceLROnPlateau adjusts learning rates, and TensorBoard provides visualization.

from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, TensorBoard
import os

# Create callbacks
callbacks = [
    EarlyStopping(
        monitor='val_loss',
        patience=15,
        restore_best_weights=True,
        verbose=1
    ),
    ModelCheckpoint(
        filepath='best_model.h5',
        monitor='val_loss',
        save_best_only=True,
        verbose=1
    ),
    ReduceLROnPlateau(
        monitor='val_loss',
        factor=0.5,
        patience=5,
        min_lr=1e-7,
        verbose=1
    ),
    TensorBoard(
        log_dir='./logs',
        histogram_freq=1
    )
]

# Train with all callbacks
history = model.fit(
    X_train, y_train,
    validation_data=(X_val, y_val),
    epochs=100,
    batch_size=32,
    callbacks=callbacks,
    verbose=1
)

This configuration saves the best model to disk (even if training continues), reduces learning rate when progress stalls (before early stopping triggers), logs metrics for TensorBoard visualization, and stops training when appropriate.

The callback execution order matters: ReduceLROnPlateau should have shorter patience than EarlyStopping so the learning rate adjusts before training stops. ModelCheckpoint ensures you have the best model saved regardless of when training ends.

Early stopping is essential for efficient, effective deep learning. Use TensorFlow’s built-in callback for standard scenarios, customize when needed, and combine with other callbacks for production-grade training pipelines. The key is finding the right balance between training thoroughness and computational efficiency for your specific use case.

Liked this? There's more.

Every week: one practical technique, explained simply, with code you can use immediately.