Stop Keras Training at the Best Epoch with restore_best_weights

The problem: training past the best epoch

You compile a Keras model and call model.fit() for 100 epochs. Validation loss improves for the first dozen epochs, then stops moving. The model continues training anyway, begins to overfit, and you are left hunting through logs to identify which epoch actually performed best before you save a checkpoint.

The useful takeaway: tf.keras.callbacks.EarlyStopping with restore_best_weights=True halts training when a validation metric stops improving and restores the weights from the best recorded epoch, so you do not have to pick the final epoch by hand.

How EarlyStopping decides when to stop

The callback watches a metric in the training logs, most commonly val_loss. Validation data is a held-out split used to estimate generalization during training. Patience is the number of epochs with no improvement after which training stops. When restore_best_weights=True, Keras copies the weights from the epoch with the best monitored value back into the model at the end.

It requires validation data or a validation split, because val_loss must exist in the logs. The metric name must match exactly; custom metrics appear under whatever name you gave them during compilation.

Worked example

Run this in a Python script or Jupyter notebook where TensorFlow 2.x is already installed. No special permissions are required. Replace the placeholder variables with your own data and model.

import tensorflow as tf
from tensorflow import keras

# Placeholders: replace with your compiled model and datasets
model = keras.Sequential([...])
train_features, train_labels = ..., ...
val_features, val_labels = ..., ...

early_stop = keras.callbacks.EarlyStopping(
    monitor="val_loss",
    patience=3,
    restore_best_weights=True,
    verbose=1,
)

history = model.fit(
    train_features,
    train_labels,
    validation_data=(val_features, val_labels),
    epochs=100,
    callbacks=[early_stop],
)

# Diagnostic: check when training stopped
print("Stopped at epoch:", early_stop.stopped_epoch)

Expected checks: early_stop.stopped_epoch reports the epoch at which the patience window expired. Training should stop roughly patience epochs after the best validation loss was recorded. Evaluate on a held-out test set and confirm the result reflects the best validation performance rather than the final epoch.

Risks: omitting validation_data means val_loss is never logged, so the callback cannot detect improvement. Setting patience too high wastes compute on unproductive epochs after the best weights have already passed.

Limitations and trade-offs

• In-memory only: restore_best_weights=True restores weights into the current model object but does not write a file to disk. Pair it with ModelCheckpoint if you need a persistent saved model.
• Exact metric name: the callback only triggers if the monitored key exists in the logs under the exact name, such as val_loss.
• Best epoch may be early: a long patience setting means many epochs run after the peak, increasing training time without improving the restored model.

Practical verification

1. Fit the model with the callback and log history.history["val_loss"].
2. Identify the minimum validation loss and its epoch.
3. Confirm early_stop.stopped_epoch is approximately that epoch plus patience.
4. Run model.evaluate(test_features, test_labels) and verify the test metric is consistent with the best validation value, not the last epoch's value.

When to use it

Add EarlyStopping with restore_best_weights=True whenever you have validation data and want to stop training automatically instead of manually selecting an epoch. Use ModelCheckpoint alongside it when you also need a permanent artifact on disk.