Custom methods and training configuration in subclass model not saved by ModelCheckpoint callback

[lilianabrandao]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

binary

TensorFlow version

tf 2.12.0

Custom code

Yes

OS platform and distribution

Google Colab / Windows 11

Mobile device

No response

Python version

3.8.18

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

Hi,

I'm following the tutorial in https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit to implement a subclassed model (end-to-end GAN). For that:

• I have overridden certain methods (train_step, save).
• I am using the ModelCheckpoint callback to save the model (generator).

The training and model saving run smoothly, but when I load the trained model:

• The overridden methods are not accessible; instead, the methods from the parent class (tf.keras.Model) are loaded.
  
• The training configuration is also not available.
  

I attempted the fix suggested in https://github.com/tensorflow/tensorflow/issues/62450 (and the related suggestions mentioned there), but without success.

Standalone code to reproduce the issue

https://colab.research.google.com/drive/19tc4EY_wI66ZZBeLlhs0rqIyfQIKENhd#scrollTo=OM_MBvzrtm_z

Relevant log output

No response

[tilakrayal]

@lilianabrandao,
I tried to execute the mentioned code with the latest TensorFlow v2.17 which contains Keras3.0 and observed that the compile error which you are facing is unavailable. And also the shape mismatch error is occured due to input difference. Kindly find the gist of it here. Also this issue is more related to Keras please create the issue in Keras-team/keras repo. Thank you!

[lilianabrandao]

Thanks for the gist, @tilakrayal, but I need to address this issue under TensorFlow 2.12.0.

After some investigation, I realized that since the subclassed model is just a wrapper for training two functional models, I need to:

1. Explicitly compile the functional models inside the wrapper to save their training configurations.
2. Override the call function in the wrapper model and explicitly invoke it before saving the subclassed model.

I hope this helps anyone facing the same issue in the future.

[google-ml-butler]

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