How to resume Keras training from a checkpoint

Interrupted Keras training should restart from a recent training state instead of starting from epoch zero. keras.callbacks.BackupAndRestore writes a temporary recovery checkpoint during model.fit() and restores it when the same training job starts again.

This callback is for fault tolerance during a training run. It is different from ModelCheckpoint, which saves a best or latest model artifact for promotion, evaluation, or later inference.

Use the same model architecture, compile settings, fit arguments, and backup directory when restarting. Changing those inputs can make the temporary recovery checkpoint invalid.

1. Install Keras and the backend package used by the training job.

   $ python -m pip install --upgrade keras jax

2. Create resume_training_checkpoint.py.

   resume_training_checkpoint.py

   import os
   import shutil
   from pathlib import Path
    
   os.environ["KERAS_BACKEND"] = "jax"
    
   import keras
   import numpy as np
    
    
   class InterruptAtEpoch(keras.callbacks.Callback):
       def __init__(self, epoch_to_stop):
           super().__init__()
           self.epoch_to_stop = epoch_to_stop
    
       def on_epoch_begin(self, epoch, logs=None):
           if epoch == self.epoch_to_stop:
               raise RuntimeError(f"simulated interruption at epoch {epoch}")
    
    
   def build_model():
       model = keras.Sequential(
           [
               keras.Input(shape=(4,), name="features"),
               keras.layers.Dense(8, activation="relu"),
               keras.layers.Dense(1, activation="sigmoid"),
           ]
       )
       model.compile(
           optimizer=keras.optimizers.Adam(learning_rate=0.03),
           loss=keras.losses.BinaryCrossentropy(),
           metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
       )
       return model
    
    
   keras.utils.set_random_seed(17)
   backup_dir = Path("training_backup")
   if backup_dir.exists():
       shutil.rmtree(backup_dir)
    
   x_train = np.linspace(0.0, 1.0, 96, dtype="float32").reshape(24, 4)
   y_train = ((x_train[:, 0] + x_train[:, 2]) > 0.75).astype("float32")
    
   model = build_model()
   callback = keras.callbacks.BackupAndRestore(backup_dir=backup_dir)
    
   try:
       model.fit(
           x_train,
           y_train,
           epochs=6,
           batch_size=6,
           callbacks=[callback, InterruptAtEpoch(3)],
           verbose=0,
       )
   except RuntimeError as exc:
       print(f"first run stopped: {exc}")
    
   resumed_model = build_model()
   resume_callback = keras.callbacks.BackupAndRestore(backup_dir=backup_dir)
   history = resumed_model.fit(
       x_train,
       y_train,
       epochs=6,
       batch_size=6,
       callbacks=[resume_callback],
       verbose=0,
   )
    
   print(f"backend: {keras.backend.backend()}")
   print(f"backup directory: {backup_dir}")
   print("interrupted before epoch: 3")
   print(f"resumed epochs completed: {len(history.history['loss'])}")
   print(f"final accuracy: {history.history['accuracy'][-1]:.4f}")
   print("training target reached: epoch 6")

   The script rebuilds the model before the second fit() call to prove that the state is restored from the backup directory, not merely kept in memory.

3. Run the script.

   $ python resume_training_checkpoint.py
   first run stopped: simulated interruption at epoch 3
   backend: jax
   backup directory: training_backup
   interrupted before epoch: 3
   resumed epochs completed: 3
   final accuracy: 0.7083
   training target reached: epoch 6

   Confirm that the second run completed only the remaining epochs because BackupAndRestore restored the saved epoch state.

4. Keep the backup directory private to one training job.

   Do not reuse the same backup_dir for another model, callback, or experiment. Use a separate checkpoint or model-save path for durable artifacts.