What Is Mixed Precision Training in TensorFlow?

Mixed precision is a training technique that combines 16-bit floating-point operations (float16 or bfloat16) with 32-bit operations (float32) to speed up deep learning training while preserving numerical stability. According to TensorFlow’s own mixed precision guide, the technique improves performance by more than 3x on modern NVIDIA GPUs with compute capability 7.0 or higher, by around 60% on Cloud TPUs, and by more than 2x on recent Intel CPUs — by running most operations in float16 (or bfloat16 on TPUs) while keeping the loss and output layer in float32.

When you are training large Machine Learning models, you want to minimise the training time. In TensorFlow, you can enable mixed precision model training, which gives a real performance improvement because it runs most operations at 16 bits (float16) instead of full single precision (float32). Google TPUs and NVIDIA GPUs have dedicated hardware — Tensor Cores, on the NVIDIA side — built to run 16-bit operations faster than 32-bit ones, see Mixed precision. For the official API reference covering tf.keras.mixed_precision.set_global_policy, the mixed_float16 policy, and the full list of supported policies, see TensorFlow’s mixed_precision documentation. The performance gain comes from two things: float16 values take half the memory of float32, so more data fits through memory bandwidth per second, and Tensor Cores execute float16 matrix multiplications and convolutions directly in hardware. In this post, I will show you the main steps for setting up mixed-precision training in TensorFlow, and then outline the data types behind them.

Mixed Precision in TensorFlow: Setup and the mixed_float16 Policy

Higher precision floats occupy 32 or 64 bits of memory and need more computational resources to process. So it makes sense to use lower-precision numbers such as float16 where you can — modern accelerators have dedicated hardware for 16-bit computations, and TensorFlow’s mixed precision guide explains how to combine 16-bit and 32-bit float operations for more efficient model training. I will walk you through a working example, following that guide together with the TensorFlow Developer Certificate in 2022: Zero to Mastery course material and the “Mixed precision” Colab notebook from TensorFlow.

Following the TensorFlow Developer Certificate in 2022: Zero to Mastery course, you implement mixed-precision TensorFlow models in two steps:

1. Activate mixed precision with mixed_precision.set_global_policy("mixed_float16").
2. Set the output layer to use float32.

Next, I will show you in detail how to confirm mixed precision is active, and how to check your hardware.

Keras Mixed Precision: Setting mixed_float16 with tf.keras.mixed_precision.set_global_policy()

I run the following code to check GPU and TPU availability and set the mixed-precision policy.

import tensorflow as tf

try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection
    policyConfig = 'mixed_bfloat16'
except ValueError:
    policyConfig = 'mixed_float16'

policy = tf.keras.mixed_precision.Policy(policyConfig)
tf.keras.mixed_precision.set_global_policy(policy)

The warning this prints, its fix and the data types behind it follow below.

Computer Data Types: Bits, Bytes, and Numeric Ranges

🔒 Subscribe to keep reading.

Floating-point Format: float16 vs float32 vs float64

🔒 Subscribe to keep reading.

Final Thoughts: When to Use Mixed Precision Training

🔒 Subscribe to keep reading.

References

🔒 Subscribe to keep reading.