TopolearnLong jobs are preempted and crash; checkpoints let you resume exactly where you stopped.
A checkpoint stores variable values so training can resume. Unlike SavedModel, it contains no graphs: you need the code that builds the objects to restore into.
tf.train.Checkpoint(**objects) saves any trackable objects you pass: models, optimizers, tf.Variables such as step counters, even datasets' iterators. Matching is object-based: variables are found by their path in the object graph (model/layer/kernel), not by global names, so reordering code does not break restores.
tf.train.CheckpointManager(ckpt, directory, max_to_keep=3) adds numbering, keeps the last N checkpoints, and exposes latest_checkpoint. A resumable loop:
ckpt.restore(manager.latest_checkpoint) # no-op if Nonethen train and call manager.save() every N steps.
Restoration can be deferred: if a variable does not exist yet (e.g. optimizer slots created on the first step), its value is applied when it is created. Use status.expect_partial() to silence warnings when you deliberately restore only part of the state, such as weights for inference.
With fit(), the keras.callbacks.ModelCheckpoint callback (or BackupAndRestore for fault tolerance) is simpler and saves Keras-format files.
Example
import tensorflow as tf
model = tf.keras.Sequential([tf.keras.Input((3,)), tf.keras.layers.Dense(1)])
step = tf.Variable(0, dtype=tf.int64)
ckpt = tf.train.Checkpoint(model=model, step=step)
manager = tf.train.CheckpointManager(ckpt, "ckpts", max_to_keep=3)
for _ in range(5):
step.assign_add(1) # ... training would happen here ...
manager.save()
print(len(manager.checkpoints)) # 3 (older ones deleted)
print(manager.latest_checkpoint) # ckpts/ckpt-5
step.assign(0)
ckpt.restore(manager.latest_checkpoint)
print(step.numpy()) # 5Remember
- Checkpoints hold variable values, matched by object path
- CheckpointManager handles numbering and retention
- Use
expect_partial()when restoring only part of the state
Watch out for
- Restoring before the model is built and assuming values are already loaded; restoration may be deferred until variables exist.