How to resume Keras training from a checkpoint

Long-running model training can lose hours of progress when a process exits between epochs. Keras can keep a temporary recovery checkpoint that lets a replacement process continue the same model.fit() run from its last saved epoch.

The BackupAndRestore callback stores the model weights, optimizer state, and completed epoch number in one dedicated directory. Its default save_freq=“epoch” setting replaces that recovery state after each epoch, limiting the amount of repeated training after an interruption.

The restarted process must rebuild the same model and use the same compile settings, fit arguments, and backup_dir. A successful fit removes the temporary backup by default, so use ModelCheckpoint separately when a durable model artifact must remain after training.

1. Create resume_training_checkpoint.py with the backend selection, constants, and deterministic training data.

   resume_training_checkpoint.py

   import argparse
   import os
   from pathlib import Path
    
   os.environ.setdefault("KERAS_BACKEND", "jax")
    
   import keras
   import numpy as np
    
    
   BACKUP_DIR = Path("training_backup")
   EPOCHS = 6
    
   keras.utils.set_random_seed(17)
   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")

2. Add the model factory below the training data.

   resume_training_checkpoint.py

   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

   The replacement process must construct a model compatible with the state stored in training_backup.

3. Add the simulated interruption callback below build_model().

   resume_training_checkpoint.py

   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 before epoch {epoch + 1}")

   The callback interrupts only the first demonstration run; a real replacement process starts without it.

4. Append the command-line option and callback selection below the interruption class.

   resume_training_checkpoint.py

   parser = argparse.ArgumentParser()
   parser.add_argument("--interrupt", action="store_true")
   args = parser.parse_args()
    
   callbacks = [keras.callbacks.BackupAndRestore(backup_dir=BACKUP_DIR)]
   if args.interrupt:
       callbacks.append(InterruptAtEpoch(epoch_to_stop=3))
    
   model = build_model()

   A dedicated backup_dir prevents another model or callback from overwriting the only recovery state.

5. Append the training call with its interruption report below the callback setup.

   resume_training_checkpoint.py

   try:
       history = model.fit(
           x_train,
           y_train,
           epochs=EPOCHS,
           batch_size=6,
           callbacks=callbacks,
           verbose=0,
       )
   except RuntimeError as exc:
       print(exc)
       print(f"backup directory exists: {BACKUP_DIR.is_dir()}")
       raise SystemExit(75)

6. Complete the script with restored-epoch and backup-removal checks.

   resume_training_checkpoint.py

   completed_epochs = [epoch + 1 for epoch in history.epoch]
   backup_removed = not BACKUP_DIR.exists()
    
   assert completed_epochs == [4, 5, 6]
   assert backup_removed
    
   print(f"completed epochs in resumed run: {completed_epochs}")
   print(f"temporary backup removed after completion: {backup_removed}")

   The epoch assertion fails if the second process starts from epoch 1 instead of restoring the three completed epochs.

7. Run the script with --interrupt to leave recovery state after epoch 3.

   $ python resume_training_checkpoint.py --interrupt
   simulated interruption before epoch 4
   backup directory exists: True

8. Verify restored epochs 4 through 6 by rerunning the script without --interrupt.

   $ python resume_training_checkpoint.py
   completed epochs in resumed run: [4, 5, 6]
   temporary backup removed after completion: True