tensorflow
deep learning
early stopping
machine learning
model training

how to implement early stopping in tensorflow

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Early stopping is one of the simplest ways to control overfitting in TensorFlow training loops. Instead of training for a fixed number of epochs and hoping you guessed well, you monitor a validation metric and stop once the model stops improving for long enough.

The standard Keras callback

In TensorFlow, the normal implementation is tf.keras.callbacks.EarlyStopping. You attach it to model.fit and tell it which metric to watch.

python
1import tensorflow as tf
2from tensorflow import keras
3
4(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
5x_train = x_train.astype("float32") / 255.0
6x_test = x_test.astype("float32") / 255.0
7
8model = keras.Sequential([
9    keras.layers.Flatten(input_shape=(28, 28)),
10    keras.layers.Dense(128, activation="relu"),
11    keras.layers.Dense(10, activation="softmax"),
12])
13
14model.compile(
15    optimizer="adam",
16    loss="sparse_categorical_crossentropy",
17    metrics=["accuracy"],
18)
19
20early_stopping = keras.callbacks.EarlyStopping(
21    monitor="val_loss",
22    patience=3,
23    restore_best_weights=True,
24)
25
26history = model.fit(
27    x_train,
28    y_train,
29    validation_split=0.2,
30    epochs=50,
31    batch_size=128,
32    callbacks=[early_stopping],
33)

This stops training when val_loss fails to improve for three consecutive epochs and restores the weights from the best epoch.

Choosing the right metric

The monitored metric should match the real goal:

  • use val_loss for most regression and classification problems
  • use val_accuracy only when accuracy is the metric you genuinely care about
  • use a custom validation metric if the default metrics do not reflect business quality

In many cases, val_loss is a better default than accuracy because it reacts earlier to overfitting.

Why restore_best_weights matters

Without restore_best_weights=True, training stops at the last epoch that exceeded patience, not at the best epoch. That often leaves the model slightly worse than the best point seen during training.

If you want the best checkpoint even when training runs longer than necessary, combine early stopping with model checkpointing:

python
1checkpoint = keras.callbacks.ModelCheckpoint(
2    "best_model.keras",
3    monitor="val_loss",
4    save_best_only=True,
5)
6
7early_stopping = keras.callbacks.EarlyStopping(
8    monitor="val_loss",
9    patience=3,
10    restore_best_weights=True,
11)
12
13model.fit(
14    x_train,
15    y_train,
16    validation_split=0.2,
17    epochs=50,
18    callbacks=[checkpoint, early_stopping],
19)

This gives you both a good stopping rule and a saved copy of the best model.

What patience actually does

patience is not the number of bad batches or noisy updates. It is the number of epochs with no meaningful improvement. A patience of 0 is usually too aggressive because validation curves often bounce slightly. A small value such as 2 through 5 is a common starting point.

TensorFlow also supports min_delta, which defines how large an improvement must be before it counts:

python
1early_stopping = keras.callbacks.EarlyStopping(
2    monitor="val_loss",
3    patience=4,
4    min_delta=0.001,
5    restore_best_weights=True,
6)

This prevents tiny fluctuations from resetting the patience counter.

Early stopping in custom training loops

If you are not using model.fit, you can still implement the same idea manually: track the best validation metric, count epochs without improvement, and break once the counter exceeds your patience value. The logic is simple; EarlyStopping just packages it cleanly for Keras workflows.

Common Pitfalls

  • Monitoring training loss instead of validation loss, which defeats the point of early stopping.
  • Forgetting restore_best_weights, then keeping weights from a later, worse epoch.
  • Using no validation data at all, which leaves nothing meaningful to monitor.
  • Setting patience too low and stopping during normal metric noise.
  • Treating early stopping as a substitute for good data splits and sensible model design.

Summary

  • Use tf.keras.callbacks.EarlyStopping with a validation metric.
  • 'val_loss is usually the safest default metric to monitor.'
  • 'restore_best_weights=True is often what you actually want.'
  • Tune patience and min_delta based on how noisy the validation curve is.
  • Combine early stopping with checkpointing when you want the best saved model as well as a clean stopping rule.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.