Saving a keras model with the data augmentation layer

[ariG23498]

I am trying to add a keras model to the hub using the push_to_hub_keras API

https://github.com/huggingface/huggingface_hub/blob/4ecdf28db59ad93569378fe79abd87ec6d5ed071/src/huggingface_hub/keras_mixin.py#L60

It seems to error out due to the inclusion of an augmentation layer inside of the model. A stackoverflow thread outlining the issue.

One can reproduce the error using this coalb notebook

[merveenoyan]

@ariG23498 hello 👋🏼 Can you downgrade TF to 2.6 and try again? should work. for people who come across this: TF 2.7 has a bug with serialization of models with stateful objects like batch normalization or data augmentation.

[ariG23498]

Hey @merveenoyan Thanks for the quick reply. I have tried downgrading and it works.

But I do not think this can be a permanent solution to the problem. There are some APIs that 2.7 provides which are either in experimental stage or non existant in 2.6. I would need some more time to hack into the API to have a proper solution, but if anyone wants to look into this or already has a solution with 2.7 support, that would be great.

[merveenoyan]

@ariG23498 this is indeed a problem, team at TF has solved it in nightly AFAIK, will probably be gone soon. Under the hood, push_to_hub_keras is calling tf.keras.models.save_model and we’d have to change it’s behavior to solve it. I’ll look into that.