How to use legacy Keras 2 with TensorFlow

TensorFlow projects that still depend on Keras 2 behavior can break after a TensorFlow upgrade because tf.keras now loads Keras 3 by default in newer releases. The legacy package keeps older tf.keras code running while a migration to standalone Keras 3 is staged and tested.

The fallback uses the separate tf-keras PyPI package, imported as tf_keras, and the TF_USE_LEGACY_KERAS=1 process switch. TensorFlow checks that switch as it loads, so the variable must be present before the first TensorFlow import in each shell, notebook kernel, worker, or application server process.

Treat the switch as compatibility mode, not a replacement for migration testing. Keep TensorFlow and tf-keras on matching minor releases in locked environments, restart long-running processes after changing the launch environment, and confirm the resolved module path before running production jobs.

1. Install TensorFlow and tf-keras in the project environment.

   $ python -m pip install --upgrade "tensorflow>=2.16" tf-keras
   Successfully installed tensorflow-2.21.0 tf-keras-2.21.0

   For production lock files, pin TensorFlow and tf-keras to the same minor release, such as tensorflow~=2.21.0 and tf-keras~=2.21.0.
   Related: How to install Keras with pip

2. Confirm that the upgraded environment uses Keras 3 by default.

   $ python - <<'PY'
   import tensorflow as tf
   
   print(f"default_dense_module={tf.keras.layers.Dense.__module__}")
   PY
   default_dense_module=keras.src.layers.core.dense

   If the output already starts with tf_keras, a parent shell, notebook launcher, or service unit is already setting TF_USE_LEGACY_KERAS.

3. Export the legacy switch in the shell that launches the TensorFlow process.

   $ export TF_USE_LEGACY_KERAS=1

4. Put the same switch at the top of notebooks or script entrypoints that do not inherit the shell environment.

   train.py

   import os
    
   os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
    
   import tensorflow as tf

   Set TF_USE_LEGACY_KERAS before import tensorflow, from tensorflow import keras, TensorFlow add-on packages, and project modules that import TensorFlow. A late assignment leaves the running process on the already-loaded Keras path.

5. Restart every Python process that may have imported TensorFlow before the switch was set.

   Restart notebook kernels, job workers, application servers, and interactive shells. TensorFlow does not move an already imported process from Keras 3 to tf-keras.

6. Verify that a fresh process resolves tf.keras to tf-keras.

   $ TF_USE_LEGACY_KERAS=1 python - <<'PY'
   import os
   import tensorflow as tf
   import tf_keras
   
   print(f"tensorflow={tf.__version__}")
   print(f"tf_keras={tf_keras.__version__}")
   print(f"legacy_switch={os.environ['TF_USE_LEGACY_KERAS']}")
   print(f"dense_module={tf.keras.layers.Dense.__module__}")
   PY
   tensorflow=2.21.0
   tf_keras=2.21.0
   legacy_switch=1
   dense_module=tf_keras.src.layers.core.dense

   The tf_keras.src module path proves tf.keras is using the legacy package. Code that imports keras directly still imports standalone Keras 3; change those files to import tf_keras as keras or migrate them to Keras 3.

7. Create a small smoke test that exercises tf.keras under the legacy switch.

   legacy_keras_smoke.py

   import os
    
   os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
   os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
    
   import numpy as np
   import tensorflow as tf
    
   model = tf.keras.Sequential(
       [
           tf.keras.Input(shape=(3,), name="features"),
           tf.keras.layers.Dense(2, activation="softmax", name="score"),
       ]
   )
    
   prediction = model(np.ones((1, 3), dtype="float32"), training=False)
    
   print(f"keras_module={tf.keras.layers.Dense.__module__.split('.')[0]}")
   print(f"prediction_shape={tuple(prediction.shape)}")

8. Run the smoke test through tf.keras.

   $ TF_USE_LEGACY_KERAS=1 python legacy_keras_smoke.py
   keras_module=tf_keras
   prediction_shape=(1, 2)

   keras_module=tf_keras confirms the legacy implementation handled the model call. Replace the small Sequential model with a project entry point or regression test before relying on the compatibility mode for production jobs.

9. Remove the temporary smoke-test file after the project-specific checks pass.

   $ rm legacy_keras_smoke.py