[Bug] Issue Loading Functional Subclassing Models

[DavidLandup0]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): MacOS
• TensorFlow installed from (source or binary): 2.11
• TensorFlow version (use command below): 2.11
• Python version: 3.9

Describe the problem.

When creating custom models with custom layers, to avoid passing custom_objects, we can use @keras.utils.register_keras_serializable() to register custom objects to the global object dictionary. This doesn't seem to work on Functional Subclassing models when loading/deserializing. For example:

tf_model = lib.model()
tf_model.save('effnet.h5')

tf.keras.utils.get_registered_object('lib_name>EfficientNetV2') #lib_name.yadda.EfficientNetV2

get_registered_object() returns a class correctly. When inspecting the custom objects:

tf.keras.utils.get_custom_objects()

{'Custom>Adadelta': keras.optimizers.optimizer_experimental.adadelta.Adadelta,
 'Custom>Adafactor': keras.optimizers.optimizer_experimental.adafactor.Adafactor,
 'Custom>Adagrad': keras.optimizers.optimizer_experimental.adagrad.Adagrad,
 'Custom>Adam': keras.optimizers.optimizer_experimental.adam.Adam,
 'Custom>Adamax': keras.optimizers.optimizer_experimental.adamax.Adamax,
 'Custom>AdamW': keras.optimizers.optimizer_experimental.adamw.AdamW,
 'Custom>Ftrl': keras.optimizers.optimizer_experimental.ftrl.Ftrl,
 'Custom>Nadam': keras.optimizers.optimizer_experimental.nadam.Nadam,
 'Custom>RMSprop': keras.optimizers.optimizer_experimental.rmsprop.RMSprop,
 'Custom>SGD': keras.optimizers.optimizer_experimental.sgd.SGD,
 'Custom>SharpnessAwareMinimization': keras.models.sharpness_aware_minimization.SharpnessAwareMinimization,
 'lib_name>EfficientNetV2': lib_name.yadda.EfficientNetV2 <-- Custom class is present and registered

But, when loading:

model = tf.keras.models.load_model('effnet.h5')

An error is thrown:

ValueError: Unknown layer: 'EfficientNetV2`. Please ensure you are using a `keras.utils.custom_object_scope` and that this object is included in the scope. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

It works when I pass a custom object:

model = tf.keras.models.load_model('effnet.h5'. custom_objects={'EfficientNetV2': EfficientNetV2})

But this shouldn't be required when an object is registered to be serialized. Am I using the API wrong?

Contributing.

• Do you want to contribute a PR? (yes/no): no
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing): /

Standalone code to reproduce the issue.

Provide a reproducible test case that is the bare minimum necessary to generate the problem. If possible, please share a link to Colab/Jupyter/any notebook.

Source code / logs.

Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem.

[sushreebarsa]

@DavidLandup0 In order to expedite the trouble-shooting process, please provide a code snippet to reproduce the issue reported here. Thank you!

[google-ml-butler[bot]]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler[bot]]

Closing as stale. Please reopen if you'd like to work on this further.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No