tensorflow_models must be imported before tf_keras to be able to load model

[vodp]

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• I am using the latest TensorFlow Model Garden release and TensorFlow 2. (latest and nightly suffer from the same issue)
• I am reporting the issue to the correct repository. (Model Garden official or research directory)
• I checked to make sure that this issue has not been filed already.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/tree/master/official/...

2. Describe the bug

I trained a model, and saved it using

import tf_keras
tf_keras.models.save_model('/path/to/model')

On loading, tensorflow_models must be imported before tf_keras for model loading to run without error

>>> import tensorflow_models as tfm
>>> import tf_keras
>>> m = tf_keras.models.load_model('/path/to/model')
>>>

3. Steps to reproduce

>>> import tf_keras
>>> m = tf_keras.models.load_model('/data/bincaptcha4/generated/bincaptcha4.1/train/model')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/trainer/.local/lib/python3.11/site-packages/tf_keras/src/saving/saving_api.py", line 301, in load_model
    return legacy_sm_saving_lib.load_model(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/trainer/.local/lib/python3.11/site-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/trainer/.local/lib/python3.11/site-packages/tf_keras/src/saving/legacy/saved_model/load.py", line 564, in _revive_graph_network
    inputs=[], outputs=[], name=config["name"]
                                ~~~~~~^^^^^^^^
KeyError: 'name

4. Expected behavior

tensorflow_models may not need to be imported for model loading to work, or the import order between tensorflow_models and tf_keras does not matter.

5. Additional context

Include any logs that would be helpful to diagnose the problem.

6. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04
• Mobile device name if the issue happens on a mobile device:
• TensorFlow installed from (source or binary): https://hub.docker.com/layers/tensorflow/tensorflow/nightly-gpu-jupyter/images/sha256-615c4ac8f5dc9aa005afedc7ac5b44a6f1198a2cbd6dc5ba700b13aed40dd2af, pip install tf-models-nightly
• TensorFlow version (use command below):2.20.0-dev
• Python version: 3.11.12
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: release 12.3, V12.3.107
• GPU model and memory: V100 32GB

Collect system information using our environment capture script.
https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh

Copyright (C) 2021 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.


== check pips ======================================================
numpy                         1.26.4
proto-plus                    1.26.1
protobuf                      5.29.4
tensorflow-datasets           4.9.8
tensorflow-metadata           1.17.1
tensorflow-model-optimization 0.8.0
tensorflow-text-nightly       2.20.0.dev20250316
tf_nightly                    2.20.0.dev20250423

== check for virtualenv ============================================
Not running inside a virtual environment.

== tensorflow import ===============================================
2025-05-02 15:27:46.818719: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1746199674.192845     257 gpu_device.cc:2018] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 31132 MB memory:  -> device: 0, name: Tesla V100-SXM2-32GB, pci bus id: 0000:1b:00.0, compute capability: 7.0

tf.version.VERSION = 2.20.0-dev20250423
tf.version.GIT_VERSION = v1.12.1-125261-g5c7cbdd5516
tf.version.COMPILER_VERSION = Ubuntu Clang 18.1.8 (++20240731024944+3b5b5c1ec4a3-1~exp1~20240731145000.144)

Sanity check: <tf.Tensor: shape=(1,), dtype=int32, numpy=array([1], dtype=int32)>
libcudnn not found

== env =============================================================
LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
DYLD_LIBRARY_PATH is unset

== nvidia-smi ======================================================
Fri May  2 15:28:05 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 555.42.06              Driver Version: 555.42.06      CUDA Version: 12.5     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  Tesla V100-SXM2-32GB           Off |   00000000:1B:00.0 Off |                    0 |
| N/A   29C    P0             43W /  300W |       6MiB /  32768MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

== cuda libs =======================================================
/usr/local/cuda-12.3/targets/x86_64-linux/lib/libcudart_static.a
/usr/local/cuda-12.3/targets/x86_64-linux/lib/libcudart.so.12.3.101

== tensorflow installation =========================================
tensorflow not found

== tf_nightly installation =========================================
Name: tf_nightly
Version: 2.20.0.dev20250423
Summary: TensorFlow is an open source machine learning framework for everyone.
Home-page: https://www.tensorflow.org/
Author-email: [email protected]
License: Apache 2.0
Location: /usr/local/lib/python3.11/dist-packages
Required-by: tf-models-nightly, tf_keras-nightly

== python version ==================================================
(major, minor, micro, releaselevel, serial)
(3, 11, 12, 'final', 0)

== bazel version ===================================================

You can also obtain the TensorFlow version with:

1. TensorFlow 1.0
   python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"

2. TensorFlow 2.0
   python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"
   v1.12.1-125261-g5c7cbdd5516 2.20.0-dev20250423

[bharatjetti]

Hello @vodp

I tried to run the docker image from the URL, trained a simple model, saved it and loaded it with both keras. i didn't get the error mentioned. Might be the docker image has recently updated in the specified URL, It is recommended to use latest-jupyter tagged image from docker hub. Could you retry and check the notebook for reference in tensorflow-tutorials/save_and_load.ipynb preloaded in the image. Please feel free to comeback if still facing the issue.

Thank you.