TopolearnMixed precision can roughly double training speed on modern GPUs and TPUs on every Keras backend.
Keras controls numeric precision with dtype policies. A mixed policy computes in 16-bit floats while storing variables in float32:
keras.config.set_dtype_policy("mixed_bfloat16")Every layer created afterwards gets compute_dtype bfloat16 and variable_dtype float32. Use "mixed_bfloat16" on TPUs and recent GPUs (NVIDIA Ampere and newer); use "mixed_float16" on older tensor-core GPUs.
Guidelines:
- Keep the model's output in float32 by giving the last layer
dtype="float32"(often a finalActivation("softmax", dtype="float32")). - With
mixed_float16, small gradients underflow, so aLossScaleOptimizeris needed;compile()wraps your optimizer automatically. In a custom train step, calloptimizer.scale_loss(loss)before computing gradients. - bfloat16 has float32's exponent range and does not need loss scaling.
- Individual layers can override the policy with
dtype=....
Set the policy before building the model, and reset it with keras.config.set_dtype_policy("float32") if later code should run in full precision. Measure: on CPUs and older GPUs mixed precision can be slower.
Example
import keras
from keras import layers
keras.config.set_dtype_policy("mixed_bfloat16") # or "mixed_float16" on older GPUs
model = keras.Sequential([
keras.Input((512,)),
layers.Dense(1024, activation="relu"),
layers.Dense(10),
layers.Activation("softmax", dtype="float32"), # numerically sensitive output
])
d = model.layers[0]
print(d.dtype_policy.name, d.compute_dtype, d.variable_dtype)
# mixed_bfloat16 bfloat16 float32
print(model.layers[-1].compute_dtype) # float32
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
# with "mixed_float16", compile() wraps the optimizer in a LossScaleOptimizer
keras.config.set_dtype_policy("float32") # reset the global policyRemember
- Mixed policies: 16-bit compute, float32 variables
- Keep the final output layer in float32
- float16 needs loss scaling (automatic in compile); bfloat16 does not
Watch out for
- Setting the policy after the model was built, so existing layers keep float32.