How to enable mixed precision in Keras

Keras mixed precision uses lower-precision compute for supported layers while keeping sensitive values, such as many variables and final outputs, in a safer dtype. On modern accelerators this can reduce memory use and improve training throughput without rewriting the model.

In Keras 3, the global dtype policy is set through keras.config.set_dtype_policy(). Layers created after the policy is set inherit it unless the layer receives its own dtype argument.

Use mixed_float16 for GPU-oriented float16 training and mixed_bfloat16 when the target hardware is optimized for bfloat16. Keep numerically sensitive output layers in float32 when the model returns probabilities or losses that should not be rounded to float16.

1. Install Keras and the backend package used by the training run.

   $ python -m pip install --upgrade keras jax

   The example uses the JAX backend. Use tensorflow or torch when that is the backend selected for the project.

2. Select the backend before importing Keras.

   $ export KERAS_BACKEND=jax

   Keras reads KERAS_BACKEND during import. Restart Python shells or notebook kernels after changing it.
   Related: How to set the Keras backend

3. Create enable_mixed_precision.py.

   enable_mixed_precision.py

   import os
    
   os.environ["KERAS_BACKEND"] = "jax"
    
   import keras
   import numpy as np
    
    
   keras.utils.set_random_seed(42)
   keras.config.set_dtype_policy("mixed_float16")
    
   x_train = np.linspace(0.0, 1.0, 64, dtype="float32").reshape(16, 4)
   y_train = ((x_train[:, 0] + x_train[:, 1]) > 0.7).astype("float32")
    
   model = keras.Sequential(
       [
           keras.Input(shape=(4,), name="features"),
           keras.layers.Dense(8, activation="relu", name="hidden"),
           keras.layers.Dense(1, activation="sigmoid", dtype="float32", name="score"),
       ]
   )
   model.compile(
       optimizer=keras.optimizers.Adam(learning_rate=0.02),
       loss=keras.losses.BinaryCrossentropy(),
       metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
   )
   history = model.fit(x_train, y_train, epochs=3, batch_size=4, verbose=0)
   prediction = model.predict(x_train[:2], verbose=0)
    
   print(f"backend: {keras.backend.backend()}")
   print(f"global dtype policy: {keras.config.dtype_policy().name}")
   print(f"hidden layer policy: {model.get_layer('hidden').dtype_policy.name}")
   print(f"output layer policy: {model.get_layer('score').dtype_policy.name}")
   print(f"optimizer: {model.optimizer.__class__.__name__}")
   print(f"epochs completed: {len(history.history['loss'])}")
   print(f"prediction dtype: {prediction.dtype}")
   print(f"prediction shape: {prediction.shape}")
    
   keras.config.set_dtype_policy("float32")

   The final layer explicitly uses float32 so the probability output stays in a stable dtype even though hidden layers inherit mixed_float16.

4. Run the script.

   $ python enable_mixed_precision.py
   backend: jax
   global dtype policy: mixed_float16
   hidden layer policy: mixed_float16
   output layer policy: float32
   optimizer: LossScaleOptimizer
   epochs completed: 3
   prediction dtype: float32
   prediction shape: (2, 1)

   Confirm that LossScaleOptimizer appears, which shows that Keras wrapped the optimizer for mixed-float16 training.

5. Reset the global policy in long-lived Python sessions when later code should return to normal float32 behavior.

   keras.config.set_dtype_policy("float32")