Mixed precision

Mixed precision gives large speedups on GPUs with tensor cores and on TPUs, with minimal code changes.

Mixed precision computes most layers in 16-bit floats while keeping variables in float32 for accurate updates. In Keras you enable it with a global dtype policy:

keras.mixed_precision.set_global_policy("mixed_float16")

Each layer then has compute_dtype float16 and variable_dtype float32. Use "mixed_bfloat16" on TPUs and recent GPUs; bfloat16's wide exponent range avoids underflow.

Rules of thumb:

Make the last layer's output float32, for example Activation("softmax", dtype="float32"), since softmax and losses are sensitive to precision.
With mixed_float16, gradients can underflow, so use loss scaling. model.compile wraps the optimizer in a LossScaleOptimizer automatically. In a custom loop, wrap it yourself and scale the loss with optimizer.scale_loss(loss) before computing gradients; apply_gradients unscales them.
Speedups need tensor cores (NVIDIA compute capability 7.0+) and dimensions that are multiples of 8.

On older GPUs or CPUs, mixed precision may run slower than float32, so measure.

Example

from tensorflow import keras

keras.mixed_precision.set_global_policy("mixed_float16")

model = keras.Sequential([
    keras.Input((784,)),
    keras.layers.Dense(1024, activation="relu"),
    keras.layers.Dense(10),
    keras.layers.Activation("softmax", dtype="float32"),   # float32 outputs
])
dense = model.layers[0]
print(dense.compute_dtype, dense.variable_dtype)   # float16 float32
print(model.layers[-1].compute_dtype)              # float32

# compile() wraps the optimizer in a LossScaleOptimizer for mixed_float16
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
print(type(model.optimizer).__name__)              # LossScaleOptimizer

keras.mixed_precision.set_global_policy("float32") # reset for later code

Remember

Policy mixed_float16/mixed_bfloat16: 16-bit compute, float32 variables
Keep the final activation in float32
float16 needs loss scaling; compile adds it automatically

Watch out for

Leaving the softmax in float16, which can produce NaN losses.

Check yourself

Under the mixed_float16 policy, what dtype are a Dense layer's weights stored in?

Official docs: www.tensorflow.org/guide/mixed_precision

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks