Tensorflow
model optimization
validation error
machine learning
model saving

Tensorflow save the model with smallest validation error

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

If you want TensorFlow or Keras to keep the model from the epoch with the lowest validation error, the standard solution is ModelCheckpoint with monitor='val_loss' and save_best_only=True. That tells Keras to overwrite the saved model only when the validation loss improves.

The important detail is that the model left in memory at the end of training is not automatically the best one unless you also restore the best weights. Saving the best checkpoint and ending training are separate concerns.

Save the Best Model During Training

Here is the basic pattern:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(16, activation="relu"),
6    tf.keras.layers.Dense(1)
7])
8
9model.compile(optimizer="adam", loss="mse")
10
11checkpoint = tf.keras.callbacks.ModelCheckpoint(
12    filepath="best_model.keras",
13    monitor="val_loss",
14    mode="min",
15    save_best_only=True,
16    verbose=1,
17)
18
19history = model.fit(
20    x_train,
21    y_train,
22    validation_data=(x_val, y_val),
23    epochs=50,
24    callbacks=[checkpoint],
25)

Because val_loss is something you want to minimize, mode='min' is the correct setting.

Load the Best Model After Training

If you save the best model to disk, load it explicitly before evaluation or inference.

python
best_model = tf.keras.models.load_model("best_model.keras")
loss = best_model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", loss)

This matters because the final epoch may not be the best epoch. Overfitting often means later epochs are worse on validation data than earlier ones.

Combine with EarlyStopping

A very common pattern is to pair checkpointing with early stopping.

python
1checkpoint = tf.keras.callbacks.ModelCheckpoint(
2    filepath="best_model.keras",
3    monitor="val_loss",
4    mode="min",
5    save_best_only=True,
6)
7
8early_stopping = tf.keras.callbacks.EarlyStopping(
9    monitor="val_loss",
10    mode="min",
11    patience=5,
12    restore_best_weights=True,
13)
14
15history = model.fit(
16    x_train,
17    y_train,
18    validation_data=(x_val, y_val),
19    epochs=100,
20    callbacks=[checkpoint, early_stopping],
21)

With restore_best_weights=True, the in-memory model is reset to the best validation-loss epoch when training stops. That means you may not need to reload from disk immediately, although many teams still keep the checkpoint for reproducibility.

Save Weights Only or the Whole Model

You can save either the whole model or only the weights.

python
1weights_checkpoint = tf.keras.callbacks.ModelCheckpoint(
2    filepath="best.weights.h5",
3    monitor="val_loss",
4    mode="min",
5    save_best_only=True,
6    save_weights_only=True,
7)

Saving only weights is smaller and faster, but then you must recreate the model architecture before loading them. Saving the full model is often simpler when you want an immediately reusable artifact.

Monitoring a Different Validation Metric

The same callback works for other validation metrics too. The only rule is to choose the right mode.

python
1checkpoint = tf.keras.callbacks.ModelCheckpoint(
2    filepath="best_accuracy_model.keras",
3    monitor="val_accuracy",
4    mode="max",
5    save_best_only=True,
6)

Use mode='max' for metrics that should increase, such as accuracy or AUC. Use mode='min' for loss or error metrics that should decrease.

A Practical Mental Model

Think of ModelCheckpoint as "preserve the best artifact so far" and EarlyStopping as "decide when to stop wasting epochs." They are complementary tools, not substitutes for each other.

That distinction helps avoid confusion when training logs say the best model was saved at epoch 12, but the training loop actually ends at epoch 20.

Common Pitfalls

  • Forgetting save_best_only=True, which causes every epoch to be saved instead of only the best one.
  • Using mode='max' while monitoring val_loss, which would keep worse models instead of better ones.
  • Assuming the final in-memory model is the best validation model when restore_best_weights was not enabled.
  • Monitoring a metric name that does not actually exist in the training logs.
  • Saving only weights and then forgetting that the model architecture must be rebuilt before loading them.

Summary

  • Use ModelCheckpoint with monitor='val_loss', mode='min', and save_best_only=True to save the model with the smallest validation error.
  • Reload the saved model after training unless you also restore the best weights in memory.
  • Pair checkpointing with EarlyStopping when you want training to stop after validation performance stalls.
  • Choose mode='min' for losses and mode='max' for metrics that should increase.
  • Decide early whether you want to save the full model or only the weights.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the 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.