keras-team/tf-keras

Importing `tf_keras` to use Keras 2 in TensorFlow 2.16 fails

Opened this issue · 9 comments

System information.

  • Have I written custom code (as opposed to using a stock example script provided in Keras): No.
  • OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS 14.2.1, Python 3.12.2 inside a conda environment
  • TensorFlow installed from (source or binary): binary (pip)
  • TensorFlow version (use command below): 2.16
  • Python version: 3.12.2
  • Bazel version (if compiling from source): -
  • GPU model and memory: -
    Describe the problem.

huggingface transformers currently relies on Keras 2. They mention that install the tf-keras package can make Keras 2 APIs available in TF 2.16. The getting started page mentions something similar. In a clean environment, I install the following packages:

pip install tensorflow tf_keras

Importing Keras 3.0 works as expected:

python -c "import tensorflow.keras; print(tensorflow.keras.version())"  
3.0.5

However, importing tf_keras fails with some error:

 python -c "import tf_keras"          
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/__init__.py", line 3, in <module>
    from tf_keras import __internal__
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/__internal__/__init__.py", line 3, in <module>
    from tf_keras.__internal__ import backend
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/__internal__/backend/__init__.py", line 3, in <module>
    from tf_keras.src.backend import _initialize_variables as initialize_variables
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/__init__.py", line 21, in <module>
    from tf_keras.src import applications
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/applications/__init__.py", line 18, in <module>
    from tf_keras.src.applications.convnext import ConvNeXtBase
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/applications/convnext.py", line 33, in <module>
    from tf_keras.src.engine import sequential
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/engine/sequential.py", line 24, in <module>
    from tf_keras.src.engine import functional
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/engine/functional.py", line 33, in <module>
    from tf_keras.src.engine import training as training_lib
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/engine/training.py", line 48, in <module>
    from tf_keras.src.saving import saving_api
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/saving/saving_api.py", line 25, in <module>
    from tf_keras.src.saving.legacy import save as legacy_sm_saving_lib
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/saving/legacy/save.py", line 27, in <module>
    from tf_keras.src.saving.legacy.saved_model import load_context
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/saving/legacy/saved_model/load_context.py", line 68, in <module>
    tf.__internal__.register_load_context_function(in_load_context)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'tensorflow._api.v2.compat.v2.__internal__' has no attribute 'register_load_context_function'. Did you mean: 'register_call_context_function'?

I can trigger the same error by forcing Tensorflow to use Keras 2 instead using TF_USE_LEGACY_KERAS=1 and importing tensorflow.keras. I do not think this is supposed to happen. From my understanding, installing tf_keras should enable us to use keras 2 in tf 2.16.

@MaxiBoether,
Could you please try to replace it from tensorflow import keras with import tf_keras as keras. Kindly find the gist for the reference.

Thank you!

Thanks for the swift response. This leads to the same error as outlined. Note that I tried two things: import tensorflow.keras (which is equivalent to from tensorflow import keras) which imports keras 3 as expected, but importing import tf_keras (it shouldn't matter whether I add as keras as requested since that's just the imported module name) fails. My goal is to import the legacy Keras 2 in a tf2.16 environment

Okay, I know what the issue is: By default (at least on my mac in a clean conda env), pip install tensorflow installs tf2.16, and pip install tf_keras installs tf_keras==2.15. On colab, pip install tensorflow defaults to v2.15. Since I just ran pip install tensorflow tf_keras, we run into this issue.

The version mismatch (if enforced on colab) leads to this error.

I'm having the same issue; attempting import tensorflow_hub (which imports tf_keras) fails with this error:

AttributeError: module 'tensorflow._api.v2.compat.v2.__internal__' 
has no attribute 'register_load_context_function'.
   Did you mean: 'register_call_context_function'?

This seems to be the result of some in-process refactoring of tf_keras/src/saving/legacy/saved_model that has since been completed, so presumably a later release will fix.

At the moment, the default installs of tensorflow and tensorflow_hub install tensorflow 2.16.0rc0 and tf-keras 2.15:

$ # start with new venv
$ python -m pip freeze | grep -e tensor -e keras
$ python -m pip install tensorflow keras tf-keras --no-cache-dir > /dev/null
$ python -m pip freeze | grep -e tensor -e keras
keras==3.0.5
tensorboard==2.16.2
tensorboard-data-server==0.7.2
tensorflow==2.16.0rc0
tf-keras==2.15.0

I've been trying to work around the incompatibility by reverting to a prior tensorflow release. I expect to be able to install out-of-date versions via pip like this:

$ python -m pip index versions tf-keras
WARNING: pip index is currently an experimental command.
 It may be removed/changed in a future release without prior warning.
tf-keras (2.15.0)
Available versions: 2.15.0, 2.14.1
$ python -m pip install tf-keras==2.14.1
Collecting tf-keras==2.14.1
  Using cached tf_keras-2.14.1-py3-none-any.whl.metadata (1.6 kB)
Using cached tf_keras-2.14.1-py3-none-any.whl (1.7 MB)
Installing collected packages: tf-keras
Successfully installed tf-keras-2.14.1

However, this approach is not working with tensorflow. I don't see anything obviously broken at PyPI, but I've not been able to install prior tensorflow releases via pip.

$ python -m pip install tensorflow==2.15.0
ERROR: Could not find a version that satisfies the requirement tensorflow==2.15.0 (from versions: 2.16.0rc0)
ERROR: No matching distribution found for tensorflow==2.15.0
$ python -m pip index versions tensorflow
WARNING: pip index is currently an experimental command. 
It may be removed/changed in a future release without prior warning.
ERROR: No matching distribution found for tensorflow

It appears 2.16.0rc0 is a pre-release version of tensorflow. Running the same commands on a different host (ubuntu under wsl) gives me 2.15.0.post1, which is what I would expect based on PyPI:

$ python -m pip index versions tensorflow
tensorflow (2.15.0.post1)
Available versions: 2.15.0.post1, 2.15.0, 2.14.1, 2.14.0, 2.13.1, 2.13.0, 2.12.1, 2.12.0, 2.11.1, 2.11.0, 2.10.1, 2.10.0, 2.9.3, 2.9.2, 2.9.1, 2.9.0, 2.8.4, 2.8.3, 2.8.2, 2.8.1, 2.8.0
  INSTALLED: 2.15.0.post1
  LATEST:    2.15.0.post1

Overall:

$ python -m pip freeze | grep -e tensor -e keras
keras==2.15.0
tensorboard==2.15.2
tensorboard-data-server==0.7.2
tensorflow==2.15.0.post1
tensorflow-estimator==2.15.0
tensorflow-hub==0.16.1
tensorflow-io-gcs-filesystem==0.36.0
tf-keras==2.15.0

With tensorflow 2.15.0.post1, tf-keras loads correctly under tensorflow-hub. I must have somehow misconfigured pip to pull pre-release versions on the other machine. Time for a rebuild...

When you use TensorFlow 2.16 with RC0 and testing Keras 2 path, please install pip install tf-keras~=2.16.0rc0 (will install RC2). If you use pip install tf-keras instead, it will install tf-keras==2.15 which is not compatible with TensorFlow 2.16.

The core of the issue is, in tf-keras 2.15 we didn't add a dependency to TensorFlow 2.15, so pip doesn't enforce that version compatibility. We fixed it in tf-keras 2.16, but may be we need to add a patch release for tf-keras 2.15 to avoid this.

@robertbcalhoun - Can you try this and confirm?

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

Just trying to avoid this issue becoming stale :-)

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.