Enabling Mixed Precision Training in Keras: Setup, Checks, and Recovery

What you are trying to achieve

Mixed precision training runs most layer computations in float16 while keeping model weights in float32. On GPUs with Tensor Cores (compute capability 7.0 and newer — Volta, Turing, Ampere and later), this typically increases throughput and reduces memory use, often letting you raise the batch size. On older GPUs or CPUs it usually gives little or no benefit, so the first step is confirming your hardware before changing any code.

The useful takeaway: in modern Keras this is a one-line policy change plus one output-layer adjustment, but you must verify it is actually active and that your loss curve still tracks your float32 baseline.

Prerequisites

• A recent TensorFlow/Keras installation. API names and default loss-scaling behavior have changed across versions, so check the docs for the version you have installed (pip show tensorflow).
• A supported GPU. Run nvidia-smi on the host to identify the GPU model, then confirm its compute capability is 7.0 or higher. No root permissions are needed for this check.
• A working float32 training run you can use as a baseline. Keep its loss curve and steps-per-second numbers for comparison.

Step 1: Set the global policy

Put this at the top of your training script, before building the model:

import keras

keras.mixed_precision.set_global_policy("mixed_float16")
print(keras.mixed_precision.global_policy())

Expected check: the printed policy is mixed_float16. From this point, layers compute in float16 but store variables in float32, which preserves numerical stability in weight updates.

Step 2: Keep the output layer in float32

float16 has a narrow dynamic range, and softmax or sigmoid outputs can underflow to zero. Override the dtype on the final layer so predictions and the loss are computed in full precision:

outputs = keras.layers.Dense(num_classes, activation="softmax", dtype="float32")(x)

Expected check: after building the model, inspect a layer:

layer = model.layers[1]
print(layer.compute_dtype, layer.variable_dtype)

You should see float16 for compute and float32 for variables on interior layers, and float32/float32 on the output layer.

Step 3: Handle loss scaling

Gradients in float16 can underflow to zero. Loss scaling multiplies the loss by a factor before backpropagation and divides the gradients afterward, keeping small gradients representable.

• With model.fit(): Keras wraps the optimizer automatically. Nothing extra is required.
• With a custom training loop: wrap the optimizer yourself:

opt = keras.optimizers.Adam(learning_rate=1e-3)
opt = keras.mixed_precision.LossScaleOptimizer(opt)

Then compute a scaled loss and use the optimizer's scaled gradient methods for your Keras version. Skipping this wrapper in a custom loop is the most common cause of silently stalled training.

Step 4: Verify against the baseline

Run a short mixed precision run on identical data and compare with your float32 baseline:

• Loss curve: it should track the float32 curve closely. Small deviations are normal; divergence or NaN is not.
• Throughput: steps per second should improve on Tensor Core hardware. There is no guaranteed multiplier — benchmark your own model and batch size.
• Memory: if training succeeds, try increasing the batch size; reduced memory use is one of the practical wins.

Recovery options when things go wrong

These options change training state, so revert deliberately rather than stacking fixes.

• Loss becomes NaN: enable or adjust dynamic loss scaling (the default behavior of LossScaleOptimizer in most versions), which raises the scale when gradients are finite and lowers it on overflow. If NaNs persist, revert to the float32 policy (set_global_policy("float32")) to confirm mixed precision is the cause.
• Accuracy degrades without NaN: cast numerically sensitive layers back to float32 selectively — normalization layers and large reductions are the usual suspects — by passing dtype="float32" to those layers.
• No speedup: re-check compute capability. On pre-Volta GPUs or CPU-only runs, revert to float32; mixed precision costs extra casts with no benefit there.

A note on TPUs

On TPUs, use mixed_bfloat16 instead. bfloat16 has a much wider dynamic range than float16, so loss scaling is usually unnecessary, but throughput characteristics differ — verify with the same baseline comparison.

Limitations

Some models are simply numerically sensitive and will not converge in float16 even with loss scaling. API details vary across Keras 2 (tf.keras) and Keras 3, so confirm method names against your installed version. Treat any speedup figure you have not measured yourself as untested.