How to save and restore a TensorFlow checkpoint

TensorFlow checkpoints keep training state as named object variables rather than as an inference export. They fit jobs that need to pause, resume, or inspect model weights, optimizer slots, and a training counter without packaging the model for serving.

A checkpoint follows the object graph passed to tf.train.Checkpoint. Restoring works when fresh objects use the same trackable names and compatible variable shapes, so a recreated model and optimizer can receive the saved values after their variables have been created.

Checkpoint files belong to the training workflow, while SavedModel exports belong to inference handoff. Keep the checkpoint directory private to the training job, verify the latest checkpoint path before restore, and use assert_consumed() when an exact object match matters.

1. Open a terminal in the Python environment that already imports TensorFlow.

   The demo was verified with TensorFlow 2.21.0. Use the same environment that owns the model code, because checkpoints do not store the Python class or model architecture.

2. Save the checkpoint round trip as

   checkpoint-save-restore-demo.py

   .

   checkpoint-save-restore-demo.py

   import os
   import pathlib
   import shutil
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("ERROR")
   tf.keras.utils.set_random_seed(7)
    
   checkpoint_dir = pathlib.Path("training_checkpoints")
   if checkpoint_dir.exists():
       shutil.rmtree(checkpoint_dir)
    
   features = tf.constant(
       [
           [0.2, 0.8, 0.4, 0.6],
           [0.9, 0.1, 0.7, 0.8],
           [0.1, 0.7, 0.3, 0.4],
           [0.8, 0.2, 0.6, 0.7],
           [0.3, 0.9, 0.2, 0.3],
           [0.7, 0.3, 0.8, 0.9],
       ],
       dtype=tf.float32,
   )
   labels = tf.constant([[0.0], [1.0], [0.0], [1.0], [0.0], [1.0]], dtype=tf.float32)
   probe = tf.constant([[0.15, 0.75, 0.35, 0.45], [0.85, 0.25, 0.65, 0.75]], dtype=tf.float32)
    
    
   def build_model():
       return tf.keras.Sequential(
           [
               tf.keras.layers.Input(shape=(4,)),
               tf.keras.layers.Dense(8, activation="relu"),
               tf.keras.layers.Dense(1, activation="sigmoid"),
           ],
           name="support_score",
       )
    
    
   model = build_model()
   optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
   model.compile(optimizer=optimizer, loss="binary_crossentropy")
   model.fit(features, labels, epochs=8, verbose=0)
    
   training_step = tf.Variable(int(optimizer.iterations), dtype=tf.int64)
   checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer, training_step=training_step)
   manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=3)
   save_path = manager.save()
    
   before = model(probe, training=False)
    
   restored_model = build_model()
   restored_optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
   restored_optimizer.build(restored_model.trainable_variables)
   restored_step = tf.Variable(0, dtype=tf.int64)
   restored_checkpoint = tf.train.Checkpoint(
       model=restored_model,
       optimizer=restored_optimizer,
       training_step=restored_step,
   )
   latest_path = tf.train.latest_checkpoint(checkpoint_dir)
   status = restored_checkpoint.restore(latest_path)
   status.assert_consumed()
    
   after = restored_model(probe, training=False)
   tf.debugging.assert_near(before, after)
    
   print(f"Saved checkpoint: {save_path}")
   print(f"Latest checkpoint: {latest_path}")
   print(f"Checkpoint files: {', '.join(sorted(path.name for path in checkpoint_dir.iterdir()))}")
   print(f"Restored training step: {int(restored_step.numpy())}")
   print(f"Predictions before save: {[round(float(value), 4) for value in before[:, 0]]}")
   print(f"Predictions after restore: {[round(float(value), 4) for value in after[:, 0]]}")
   print("Restore assertion: consumed")
   print("Prediction match: True")
    
   shutil.rmtree(checkpoint_dir)
   print(f"Cleaned up: {not checkpoint_dir.exists()}")

   Build or call the fresh model and optimizer before an exact restore assertion. Optimizers such as Adam create slot variables from the model variables, and those slot variables must exist before assert_consumed() can prove a full restore.

3. Run the checkpoint round trip.

   $ python3 checkpoint-save-restore-demo.py
   Saved checkpoint: training_checkpoints/ckpt-1
   Latest checkpoint: training_checkpoints/ckpt-1
   Checkpoint files: checkpoint, ckpt-1.data-00000-of-00001, ckpt-1.index
   Restored training step: 8
   Predictions before save: [0.3968, 0.646]
   Predictions after restore: [0.3968, 0.646]
   Restore assertion: consumed
   Prediction match: True
   Cleaned up: True

   Restore assertion: consumed confirms that every saved checkpoint value matched a tracked object in the restored model, optimizer, and training step. Prediction match: True confirms that the restored model variables produce the same outputs for the probe batch.

• tf.train.CheckpointManager keeps a checkpoint directory and tracks the latest checkpoint path, which is safer than hard-coding one filename after repeated training saves.
• A checkpoint stores variable values and object paths, not a standalone inference package. Use a SavedModel when another runtime needs to load the model without the original training code.
  Related: How to export a SavedModel in TensorFlow
• Use assert_consumed() for an exact restore into the same object graph. Use partial-restore APIs only when intentionally loading a subset, such as model weights without optimizer slots.