Stop Overfitting Fast: Using EarlyStopping and ModelCheckpoint in Keras

Why Overfitting Happens and How Callbacks Help

When a neural network learns too well from the training data, it starts to memorize noise instead of general patterns. The classic symptom is a low training loss but a steadily rising validation loss. Rather than waiting for the training to finish, you can stop the process early and keep the weights that performed best on the validation set. Keras offers two callbacks that make this simple: EarlyStopping and ModelCheckpoint.

1. EarlyStopping – “Stop When You’re Stuck”

EarlyStopping watches a chosen metric (e.g., val_loss) and terminates training if the metric does not improve for a specified number of epochs. The key arguments are:

• monitor – metric to watch.
• patience – epochs to wait after the last improvement.
• mode – ‘min’, ‘max’, or ‘auto’ (auto chooses based on metric name).
• restore_best_weights – if True, rolls back to the weights from the epoch with the best monitored value.

Example:

early_stop = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',
    patience=3,
    mode='min',
    restore_best_weights=True
)

Run this callback on the GPU or CPU where your training script is executed. It requires a validation split; otherwise the callback cannot determine the best weights.

2. ModelCheckpoint – “Save the Best”

ModelCheckpoint writes the model (architecture + weights or just weights) to disk whenever the monitored metric improves. Important parameters:

• filepath – path to the file (e.g., best.h5).
• monitor – metric to watch.
• save_best_only – if True, only overwrite the file when the metric improves.
• mode – same as in EarlyStopping.

Example:

checkpoint = tf.keras.callbacks.ModelCheckpoint(
    filepath='best.h5',
    monitor='val_loss',
    save_best_only=True,
    mode='min'
)

Ensure the directory is writable and that you use a unique filename to avoid accidental overwrites.

3. Combine Them – Order Matters

When both callbacks are used, the sequence in the callbacks list determines behavior:

1. Place checkpoint before early_stop so that the best weights are saved before training is halted.
2. If restore_best_weights=True is set on EarlyStopping, the final model will automatically revert to the checkpointed weights, making the separate checkpoint optional for recovery but still useful for debugging.

Full snippet:

callbacks = [checkpoint, early_stop]
history = model.fit(
    x_train, y_train,
    epochs=100,
    batch_size=128,
    validation_split=0.2,
    callbacks=callbacks
)

4. Minimal MNIST Demo

Below is a compact script that trains a small CNN on MNIST, uses EarlyStopping with patience 3, and checkpoints the best model. No output is provided; run it locally to see the training stop early.

import tensorflow as tf

# Load data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
# Normalize and reshape
x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
# Add channel dimension
x_train = tf.expand_dims(x_train, -1)
x_test = tf.expand_dims(x_test, -1)

# Simple model
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28,28,1)),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Callbacks
checkpoint = tf.keras.callbacks.ModelCheckpoint(
    filepath='best.h5',
    monitor='val_loss',
    save_best_only=True,
    mode='min'
)
early_stop = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',
    patience=3,
    mode='min',
    restore_best_weights=True
)

history = model.fit(
    x_train, y_train,
    epochs=50,
    batch_size=128,
    validation_split=0.2,
    callbacks=[checkpoint, early_stop],
    verbose=2
)

# Evaluate the final (restored) model
score = model.evaluate(x_test, y_test, verbose=0)
print(f'Test loss: {score[0]:.4f}, Test accuracy: {score[1]:.4f}')

# Load the checkpointed model explicitly
best_model = tf.keras.models.load_model('best.h5')
best_score = best_model.evaluate(x_test, y_test, verbose=0)
print(f'Best model loss: {best_score[0]:.4f}, accuracy: {best_score[1]:.4f}')

After running, you should see training stop after the validation loss plateaus for 3 epochs. The best.h5 file will contain the weights from the epoch with the lowest val_loss, and the final evaluation should match or slightly improve upon the checkpointed performance.

5. Trade‑offs & Caveats

• Noise in Validation – If the validation metric fluctuates, EarlyStopping might terminate too early. Consider increasing patience or applying a moving‑average smoothing to the metric.
• Checkpoint Overwrites – Re‑run the script with the same filepath will overwrite the previous best model. Use timestamped filenames or a dedicated directory.
• restore_best_weights Requires Validation – Without a validation split, the callback cannot determine the best epoch and will revert to the last epoch instead.
• Disk Space & I/O – Frequent checkpoints add I/O overhead. Setting save_best_only=True mitigates this by writing only when improvements occur.
• Compatibility – The examples target TensorFlow 2.x. In earlier Keras releases, the filepath format may differ (e.g., *.hdf5).

6. Quick Verification Checklist

1. Run training and watch the log: after the last improvement, the next patience epochs should produce no further decrease in val_loss before stopping.
2. Confirm best.h5 exists and its size matches the expected weight file.
3. Load the checkpoint with tf.keras.models.load_model('best.h5') and evaluate on a held‑out test set. The scores should be equal to or better than the final epoch’s scores.
4. Inspect history.history for the val_loss curve; the minimum should align with the epoch that produced the checkpoint.

When you follow these steps, you’ll efficiently prevent overfitting, save the most performant model, and keep training time reasonable.

Conclusion – Keep It Simple, Keep It Correct

EarlyStopping and ModelCheckpoint are low‑overhead, high‑impact tools in Keras. By placing the checkpoint first, setting restore_best_weights=True, and carefully choosing patience and the monitored metric, you can automate a robust training loop that stops at the right time and preserves the best weights for deployment or further analysis.