Mixed precision and multi-device training

Mixed precision policy and loss scaling, distribution strategies, how batch size scales across devices, and checkpointing a run that can be resumed.

Mixed precision

import tensorflow as tf

# enable the policy before building the model
policy = tf.keras.mixed_precision.Policy("mixed_float16")
tf.keras.mixed_precision.set_global_policy(policy)
print(policy.compute_dtype, policy.variable_dtype)   # float16 float32

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(128,)),
    tf.keras.layers.Dense(256, activation="relu"),
    tf.keras.layers.Dense(128, activation="relu"),
    # the output layer must stay float32: softmax and loss need the range
    tf.keras.layers.Dense(10, dtype="float32"),
])
model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-3),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)
# Keras wraps the optimiser in a LossScaleOptimizer automatically with this policy
print(model.optimizer)

• Variables stay float32; only the compute is float16. That keeps the small updates the optimiser makes from being rounded away.
• The final layer must be float32 for classification, because softmax and the log loss need the extra range and precision.
• Keras inserts the loss-scale optimiser automatically under the mixed policy. If you write a custom training loop you must do it yourself and skip steps with non-finite gradients.
• Expect roughly 1.5-2.5x throughput on modern GPUs (more on tensor-core hardware), and no benefit at all on a CPU.

MirroredStrategy and batch scaling

strategy = tf.distribute.MirroredStrategy()   # all visible GPUs
print("devices:", strategy.num_replicas_in_sync)

with strategy.scope():                        # build and compile inside the scope
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(128,)),
        tf.keras.layers.Dense(256, activation="relu"),
        tf.keras.layers.Dense(10, dtype="float32"),
    ])
    model.compile(
        optimizer=tf.keras.optimizers.Adam(1e-3),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=["accuracy"],
    )

# scale the batch with the number of replicas
GLOBAL_BATCH = 32 * strategy.num_replicas_in_sync
train_ds = train_ds.unbatch().batch(GLOBAL_BATCH)

# scale the learning rate roughly with the square root of the batch increase
base_lr = 1e-3 * (strategy.num_replicas_in_sync ** 0.5)

checkpoint = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(checkpoint, "ckpt", max_to_keep=3)

model.fit(train_ds, epochs=10,
          callbacks=[tf.keras.callbacks.ModelCheckpoint("best.keras",
                                                        save_best_only=True)])

Checkpointing and resuming

# a checkpoint that restores optimiser state, so a resumed run behaves like a continued one
ckpt = tf.train.Checkpoint(
    step=tf.Variable(0),
    optimizer=model.optimizer,
    model=model,
)
manager = tf.train.CheckpointManager(ckpt, "checkpoints/run1", max_to_keep=5)

ckpt.restore(manager.latest_checkpoint)
if manager.latest_checkpoint:
    print("resumed from", manager.latest_checkpoint, "at step", int(ckpt.step))

class CheckpointEachEpoch(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        ckpt.step.assign_add(1)
        path = manager.save()
        print("saved", path)

model.fit(train_ds, epochs=20, initial_epoch=int(ckpt.step),
          callbacks=[CheckpointEachEpoch()])

• Save the optimiser state, not only the weights. Restoring weights without momentum and adaptive moments restarts training in a different regime and often causes a visible loss spike.
• Save .keras for a portable exported model and a tf.train.Checkpoint for resumable training. They serve different purposes and you generally want both.
• Keep the last few checkpoints (max_to_keep). A single overwritten checkpoint turns a corrupted write into a lost run.
• Verify a checkpoint loads before you need it. A restoration path that only runs after a crash is a restoration path that has never been tested.

FAQ

Related

The tf.data pipeline in depth Debugging, profiling and export formats

Last refreshed 2026-09-18.