Saving/Loading custom methods of a subclassed model

[lilianabrandao]

Hi,

I'm following the tutorial in https://keras.io/guides/custom_train_step_in_tensorflow/#wrapping-up-an-endtoend-gan-example to implement a subclassed model (end-to-end GAN).

The subclassed model is a wrapper for training two sequential models (the generator and discriminator).

Everything works nicely, except that the custom methods defined in the subclassed model are not retained after saving and loading it, even when using the decorator @tf.keras.utils.register_keras_serializable().

How can I save the subclassed model while retaining its custom methods?

As explained in TensorFlow's guide on saving and loading models https://www.tensorflow.org/tutorials/keras/save_and_load#saving_custom_objects, this should be straightforward when using the SavedModel format ("tf").

Tensorflow/Keras version: 2.12.0 Attempted methods: Saving in the SavedModel format ("tf") and/or using the @tf.keras.utils.register_keras_serializable() decorator. Expected behavior: To save the custom methods of a subclassed model (class configuration). Standalone code to reproduce the issue: https://colab.research.google.com/drive/1o3OqG3UNIG3KywSOkhp6yPFUZPyLlLv0?usp=sharing

[sonali-kumari1]

Hi @lilianabrandao,

Thanks for reporting the issue but the drive link is not accessible for me. Can you provide accessible link ?

[lilianabrandao]

Thank you @sonali-kumari1 . Here's the accessible link: https://colab.research.google.com/drive/1o3OqG3UNIG3KywSOkhp6yPFUZPyLlLv0?usp=sharing

[fchollet]

It looks like you are using Keras 2 and the TF SavedModel format. Try with Keras 3 and the .keras file format, and it should work as you expect.