Mixed Precision Training

Two number formats, one training loop

Mixed precision does exactly what the name says: it mixes float16 (or bfloat16) for the heavy compute with float32 for the parts that need the range. The matmuls and convolutions — where almost all the FLOPs live — run in 16-bit, which is where the speedup comes from. But the master copy of the weights, and the gradient accumulation, stay in 32-bit so small updates don't vanish into rounding error.

One line to turn it on

You set a global policy with set_global_policy("mixed_float16") before you build the model, then build and train exactly as you normally would — Keras inserts the casts at the right boundaries for you. One detail earns its keep: keep your final Dense layer producing raw logits (no softmax) and pair it with from_logits=True, so the loss is computed in a numerically stable way instead of through a float16 softmax that can overflow.

The free-speed catch

The 2-3× is only free if the hardware has the matching tensor units. On a GPU without Tensor Cores (or on CPU) you get the memory savings but little or no speedup, and you've added casting overhead for nothing — so this is a measure-it, not assume-it, win.

Code

Enable mixed_float16 and keep logits in float32·python

# Enable mixed precision globally
keras.mixed_precision.set_global_policy("mixed_float16")

# Build and train normally — Keras handles the casting
model = keras.Sequential([
    keras.Input(shape=(784,)),
    layers.Dense(256, activation="relu"),  # Computes in float16
    layers.Dense(10),                        # Raw logits (no activation)
])

# Use from_logits for numerical stability with mixed precision
model.compile(
    optimizer="adam",
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)

Mixed Precision Training

Two number formats, one training loop

One line to turn it on

The free-speed catch

Code

External links

Exercise

Progress

Comments 0