Mixed Precision Training in Keras: Faster GPU Epochs Without Rewriting Your Model

The problem: your GPU is underutilized and epochs are slow

Your Keras model trains, but each epoch takes longer than it should, and GPU memory stays near the limit. You do not want to rewrite the architecture or switch frameworks just to improve throughput. The practical question is whether you can train faster on the same NVIDIA GPU without sacrificing meaningful accuracy.

What Keras mixed precision actually changes

Mixed precision means using more than one floating-point dtype during training. With tf.keras.mixed_precision.set_global_policy('mixed_float16'), Keras casts layer computations to float16 but keeps a float32 copy of trainable variables, often called master weights. Gradients flow back in float16, then Keras applies loss scaling: it multiplies the loss by a large factor so small gradients do not underflow to zero, then unscales them before the optimizer updates the master weights.

This works with the Sequential and Functional APIs, custom training loops that use Keras layers, and tf.distribute.MirroredStrategy. On TPUs the corresponding policy is mixed_bfloat16, not mixed_float16.

Enable the policy before building the model

Run this in a Python shell or Jupyter cell. No special permissions are required beyond the ability to run Python and import TensorFlow.

import tensorflow as tf

# Set this before any Keras model is constructed in the session.
tf.keras.mixed_precision.set_global_policy('mixed_float16')

# Verify the active policy.
print(tf.keras.mixed_precision.global_policy())

Expected check: the print statement should report Policy: mixed_float16. Risk: the policy is process-wide, so every model built afterward in the same session uses it unless you reset it.

Worked example with a small CNN

This example builds a tiny CIFAR-10 classifier under mixed precision. Replace batch_size and epochs with values suited to your hardware.

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, datasets

tf.keras.mixed_precision.set_global_policy('mixed_float16')

inputs = keras.Input(shape=(32, 32, 3))
x = layers.Conv2D(32, 3, activation='relu')(inputs)
x = layers.MaxPooling2D()(x)
x = layers.Flatten()(x)
outputs = layers.Dense(10, activation='softmax')(x)
model = keras.Model(inputs, outputs)

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

(x_train, y_train), _ = datasets.cifar10.load_data()
x_train = x_train.astype('float32') / 255.0

history = model.fit(x_train, y_train, batch_size=64, epochs=1, validation_split=0.1)

Where to run: a local workstation or notebook with TensorFlow 2.x and an NVIDIA GPU. After fitting, inspect history.history['loss'] for NaN or Inf values. If loss is finite and validation accuracy is close to a float32 baseline, the policy is stable for this model.

When mixed precision can go wrong

Not every layer is float16-safe. Custom ops or layers with very small activations can produce NaNs even with loss scaling. If that happens, wrap the problematic computation in tf.cast to float32, or keep that sub-model on float32. Accuracy can also shift slightly, so validate on a held-out set rather than assuming the result is identical. Mixed precision delivers the most benefit on GPUs with dedicated FP16 hardware, such as Volta, Ampere, and Ada architectures; older GPUs may not speed up and could even slow down.

Rollback

Because set_global_policy changes session state, revert to pure float32 with:

tf.keras.mixed_precision.set_global_policy('float32')

Models already constructed under mixed_float16 keep their dtype, so rebuild them after resetting the policy if you want a clean float32 comparison.

Closing

Start with the one-line policy call, confirm global_policy() reports mixed_float16, watch the loss for NaNs, and compare validation accuracy against your float32 baseline. If the numbers hold, you have a low-effort way to use your GPU more efficiently.