tf.train.Checkpoint: Save and Resume Training in TensorFlow

The Story — Save Points in a Video Game

What Goes Into a Checkpoint?

A checkpoint saves every tf.Variable reachable from the objects you give it. That is why it works with anything built on tf.Module.

import tensorflow as tf
import numpy as np, os, shutil
shutil.rmtree("/tmp/ckpt", ignore_errors=True)

rng = np.random.default_rng(0)
X = rng.normal(size=(256, 3)).astype("float32")
y = (X @ np.array([[2.0], [-1.0], [0.5]], dtype="float32") + 0.1 * rng.normal(size=(256, 1))).astype("float32")

def build():
    tf.random.set_seed(0)
    model = tf.keras.Sequential([tf.keras.Input((3,)), tf.keras.layers.Dense(1)])
    opt = tf.keras.optimizers.SGD(0.05, momentum=0.9)
    step = tf.Variable(0, dtype=tf.int64, name="step")
    return model, opt, step

model, opt, step = build()
ckpt = tf.train.Checkpoint(model=model, optimizer=opt, step=step)
manager = tf.train.CheckpointManager(ckpt, directory="/tmp/ckpt", max_to_keep=3)

@tf.function
def train_step():
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((model(X) - y) ** 2)
    grads = tape.gradient(loss, model.trainable_variables)
    opt.apply_gradients(zip(grads, model.trainable_variables))
    step.assign_add(1)
    return loss

for _ in range(50):
    loss = train_step()
    if int(step) % 10 == 0:
        path = manager.save(checkpoint_number=step)
        print(f"step {int(step):3d}  loss {float(loss):.4f}  saved {os.path.basename(path)}")

print("kept:", [os.path.basename(p) for p in manager.checkpoints])
print("files:", sorted(os.listdir("/tmp/ckpt")))

The Crash — and the Resume

We pretend the machine restarted: we throw away every Python object and build a brand-new model and optimizer. Then we restore.

del model, opt, step, ckpt, manager                         # "the machine restarted"

model, opt, step = build()                                   # fresh, random objects
ckpt = tf.train.Checkpoint(model=model, optimizer=opt, step=step)
manager = tf.train.CheckpointManager(ckpt, "/tmp/ckpt", max_to_keep=3)

print("latest:", os.path.basename(manager.latest_checkpoint))
status = ckpt.restore(manager.latest_checkpoint)
print("step after restore:", int(step))

@tf.function
def train_step():
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((model(X) - y) ** 2)
    grads = tape.gradient(loss, model.trainable_variables)
    opt.apply_gradients(zip(grads, model.trainable_variables))
    step.assign_add(1)
    return loss

for _ in range(10):
    loss = train_step()
status.assert_consumed()                                     # every saved value found a home
print(f"continued to step {int(step)}, loss {float(loss):.5f}")
print("weights:", model.layers[0].kernel.numpy().ravel().round(3))

Look Inside a Checkpoint

for name, shape in tf.train.list_variables(manager.latest_checkpoint):
    if not name.startswith("_CHECKPOINTABLE") and "OBJECT_CONFIG" not in name:
        print(f"{name:70s} {shape}")

The names are paths through your objects. Each variable is stored once, under the first path TensorFlow finds. Here the Dense kernel [3, 1] and bias [1] show up as optimizer/_trainable_variables/0 and /1, because the optimizer also points to them. _variables/2 and /3 are the momentum slots, and step is our counter. Restore matches values to variables by this object structure, not by Python variable names.

The Standard Pattern: Restore If Possible

ckpt = tf.train.Checkpoint(model=model, optimizer=opt, step=step)
manager = tf.train.CheckpointManager(ckpt, "ckpts/run1", max_to_keep=5)

ckpt.restore(manager.latest_checkpoint)          # does nothing if there is none yet
print("starting at step", int(step))

while int(step) < total_steps:
    train_step()
    if int(step) % 500 == 0:
        manager.save(checkpoint_number=step)

Try It — Crash and Resume Simulator

Checkpoint or SavedModel?

Golden Rules