Using Keras ModelCheckpoint to Save the Best Model During Training

Useful answer

The ModelCheckpoint callback in Keras automatically writes a model’s weights (or the full model) to disk whenever a monitored metric improves. This lets you recover the best‑performing model after training without manually tracking epochs.

How the callback works

During each epoch Keras calls the callback’s on_epoch_end method. The callback compares the current value of the metric you asked it to monitor (e.g., val_loss) with the best value seen so far, using the mode argument ('min' for loss, 'max' for accuracy). If the new value is better according to that mode and save_best_only=True, the callback serialises the model to the file path you supplied. If save_best_only=False it writes a file every epoch, which can quickly consume disk space.

Worked configuration

Below is a minimal, self‑contained example that you can run in a Python environment with TensorFlow 2.x installed. Adjust the paths and dataset placeholders to match your own data.

# 1. Imports
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import os

# 2. Create a tiny synthetic dataset (replace with your own)
#    Features: 1000 samples, 20 dimensions
#    Targets:  single regression value
X = np.random.randn(1000, 20).astype('float32')
y = (X[:, 0] * 2.0 + np.random.randn(1000) * 0.1).astype('float32')

# 3. Build a simple model
model = keras.Sequential([
    layers.Dense(64, activation='relu', input_shape=(20,)),
    layers.Dense(1)
])
model.compile(optimizer='adam', loss='mse', metrics=['mae'])

# 4. Ensure the output directory exists
checkpoint_dir = './tmp'
os.makedirs(checkpoint_dir, exist_ok=True)  # requires write permission

# 5. Instantiate the checkpoint callback
checkpoint = keras.callbacks.ModelCheckpoint(
    filepath=os.path.join(checkpoint_dir, 'best_model.keras'),
    monitor='val_loss',          # metric to watch
    save_best_only=True,         # only keep the best so far
    mode='min',                  # because lower loss is better
    save_format='keras'          # use the native .keras format (TF 2.x)
)

# 6. Train with validation split and the callback
history = model.fit(
    X, y,
    epochs=20,
    batch_size=32,
    validation_split=0.2,
    callbacks=[checkpoint],
    verbose=2
)

# 7. After training, verify the file exists and load it
best_path = os.path.join(checkpoint_dir, 'best_model.keras')
if os.path.exists(best_path):
    best_model = keras.models.load_model(best_path)
    # Evaluate on the same validation split used during training
    val_loss = best_model.evaluate(X[int(0.8*len(X)):], y[int(0.8*len(y)):], verbose=0)[0]
    print(f'Validation loss of the saved best model: {val_loss:.5f}')
else:
    print('Checkpoint file was not created – check the metric name and validation data.')

Configuration details

• filepath: destination for the saved model. Use an absolute path or ensure the directory exists (os.makedirs in the example).
• monitor: the exact name of a metric that appears in the logs (val_loss, val_accuracy, etc.). If the metric is not logged, the callback never triggers a save.
• save_best_only: true keeps only the best model; false creates a file each epoch (useful for debugging but risky for disk usage).
• mode: 'min' for metrics where lower is better (loss, MAE), 'max' for higher‑is‑better metrics (accuracy, AUC).
• save_format: In TensorFlow 2.x prefer 'keras' (produces a .keras archive) or 'tf' for SavedModel format. The legacy HDF5 (.h5) format is still supported but has size limits and may cause version‑compatibility issues.

Limits and common mistakes

• Missing validation data: If you forget to pass validation_data or validation_split, the metric you monitor (e.g., val_loss) never appears in the logs, so no checkpoint is written.
• Incorrect metric name: Typos or case mismatches (Val_Loss vs val_loss) lead to silent failures. Print history.history.keys() after training to see what was logged.
• Directory not writable: Supplying a path like '/protected/folder/model.keras' without appropriate permissions raises FileNotFoundError or PermissionError. Create the folder beforehand or choose a writable location.
• Disk‑full scenario: Setting save_best_only=False with many epochs or large models can fill the filesystem quickly. Monitor disk usage or switch to the best‑only mode.
• Metric not improving: If the model’s validation loss plateaus or worsens, the callback will not update the file after the initial best epoch. This is expected behaviour; you may want to add EarlyStopping to halt training when no improvement occurs for a set number of epochs.

Practical verification steps

1. Run the script above in a terminal or notebook where you have write access to ./tmp.
2. After training finishes, check that ./tmp/best_model.keras exists (ls -l ./tmp).
3. Load the model with keras.models.load_model and evaluate it on the validation split used during training. The reported loss should be equal to or lower than the loss of the final epoch shown in history.history['val_loss'].
4. Optionally, compare the file size of the checkpoint with a per‑epoch HDF5 save (save_best_only=False) to see the storage difference.

When to use ModelCheckpoint

Use this callback whenever you want a reliable way to retrieve the model that achieved the best validation performance during a training run, especially when training for many epochs or when you plan to deploy the model later. Combine it with EarlyStopping or ReduceLROnPlateau for a robust training pipeline.