How to enable mixed precision in TensorFlow

Mixed precision lets TensorFlow and Keras run supported model math in 16-bit types while keeping variables and selected numerically sensitive operations in float32. It is most useful when training on accelerator hardware where matrix multiplication, convolution, or memory bandwidth limits the batch throughput.

The setting is a process-local Keras dtype policy, so it affects layers created after the policy is set. Use mixed_float16 for NVIDIA GPU training, and use mixed_bfloat16 for TPU or supported CPU paths where bfloat16 is the preferred lower-precision type.

Set the policy before building the model, then keep the final prediction or loss-facing output in float32 when the output feeds a loss, probability, or metric calculation. Keras model.fit() handles dynamic loss scaling for mixed_float16, while custom GradientTape loops need explicit loss-scaling handling.

1. Activate the Python environment that runs the training job.

   $ source ~/venvs/tf-gpu/bin/activate
   (tf-gpu) $

   Use conda activate <name> instead when the project uses Conda.
   Related: How to create a virtual environment for TensorFlow
   Related: How to create a Conda environment for TensorFlow

2. Save a mixed precision smoke test as

   train_mixed_precision.py

   .

   train_mixed_precision.py

   import os
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("ERROR")
   tf.keras.utils.set_random_seed(42)
    
   mixed_precision = tf.keras.mixed_precision
   mixed_precision.set_global_policy("mixed_float16")
    
   inputs = tf.keras.Input(shape=(4,), name="features")
   x = tf.keras.layers.Dense(16, activation="relu", name="hidden")(inputs)
   outputs = tf.keras.layers.Dense(1, activation="sigmoid", dtype="float32", name="score")(x)
   model = tf.keras.Model(inputs, outputs)
    
   model.compile(
       optimizer="adam",
       loss="binary_crossentropy",
       metrics=["accuracy"],
   )
    
   features = tf.random.normal((32, 4))
   labels = tf.cast(tf.reduce_sum(features, axis=1, keepdims=True) > 0, tf.float32)
   history = model.fit(features, labels, epochs=1, batch_size=8, verbose=0)
    
   print(f"TensorFlow {tf.__version__}")
   print(f"global_policy={mixed_precision.global_policy().name}")
   print(f"hidden_compute_dtype={model.get_layer('hidden').compute_dtype}")
   print(f"hidden_variable_dtype={model.get_layer('hidden').variable_dtype}")
   print(f"output_compute_dtype={model.get_layer('score').compute_dtype}")
   print(f"prediction_dtype={model(features[:2]).dtype.name}")
   print(f"loss={history.history['loss'][-1]:.4f}")
   print(f"accuracy={history.history['accuracy'][-1]:.4f}")

   The global policy call must run before creating layers. The hidden layer inherits mixed_float16, while the final layer overrides the compute dtype to float32.

3. Run the smoke test and confirm the policy, layer dtypes, output dtype, and training result.

   (tf-gpu) $ python train_mixed_precision.py
   TensorFlow 2.21.0
   global_policy=mixed_float16
   hidden_compute_dtype=float16
   hidden_variable_dtype=float32
   output_compute_dtype=float32
   prediction_dtype=float32
   loss=0.8104
   accuracy=0.4688

   A CPU-only process can validate the API path, but mixed_float16 normally speeds up training only on recent NVIDIA GPUs. Use mixed_bfloat16 instead when the target accelerator is TPU or a supported CPU bfloat16 path.

4. Put the global policy call near the top of the real training entry point.

   import tensorflow as tf
    
   mixed_precision = tf.keras.mixed_precision
   mixed_precision.set_global_policy("mixed_float16")

   Do not set the policy after creating the model, optimizer, or layers. Existing layers keep the policy they were created with.

5. Keep the final model output in float32 when it feeds a loss or probability output.

   inputs = tf.keras.Input(shape=(num_features,))
   x = tf.keras.layers.Dense(256, activation="relu")(inputs)
   x = tf.keras.layers.Dense(128, activation="relu")(x)
   outputs = tf.keras.layers.Dense(
       num_classes,
       activation="softmax",
       dtype="float32",
       name="predictions",
   )(x)
   model = tf.keras.Model(inputs, outputs)

   Intermediate layers can keep the mixed policy. The output override avoids passing float16 predictions into the loss or downstream probability handling.

6. Compile and train the model with Keras model.fit().

   model.compile(
       optimizer="adam",
       loss="sparse_categorical_crossentropy",
       metrics=["accuracy"],
   )
    
   model.fit(train_ds, validation_data=val_ds, epochs=10)

   Keras applies dynamic loss scaling automatically for mixed_float16 during model.fit(). Custom tf.GradientTape loops that apply optimizer steps directly should use tf.keras.mixed_precision.LossScaleOptimizer.
   Related: How to run a custom training loop in TensorFlow

7. Remove the smoke-test file after the project training path reports the expected policy and dtypes.

   (tf-gpu) $ rm train_mixed_precision.py