tf-keras icon indicating copy to clipboard operation
tf-keras copied to clipboard

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

Open MaxiBoether opened this issue 1 year ago • 11 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 avatar Mar 01 '24 10:03 MaxiBoether

@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!

tilakrayal avatar Mar 01 '24 10:03 tilakrayal

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

MaxiBoether avatar Mar 01 '24 10:03 MaxiBoether

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.

MaxiBoether avatar Mar 01 '24 10:03 MaxiBoether

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

robertbcalhoun avatar Mar 03 '24 17:03 robertbcalhoun

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...

robertbcalhoun avatar Mar 03 '24 17:03 robertbcalhoun

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?

sampathweb avatar Mar 05 '24 17:03 sampathweb

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.

github-actions[bot] avatar Mar 21 '24 01:03 github-actions[bot]

Just trying to avoid this issue becoming stale :-)

MaxiBoether avatar Mar 21 '24 15:03 MaxiBoether

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.

github-actions[bot] avatar Apr 05 '24 01:04 github-actions[bot]

@MaxiBoether, I tried with the latest tensorflow 2.19 and tf-keras, and observed I was able to import without any issue/error. Kindly find the gist of it here and also by default tensorflow 2.19 contains keras3.0, if you want to use keras2, try to use pip install tf-keras and then use import tf_keras as keras. Thank you!

tilakrayal avatar Apr 29 '25 06:04 tilakrayal

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.

github-actions[bot] avatar May 14 '25 02:05 github-actions[bot]