Some custom objects are not being serialized with push_to_hub_keras

Self-contained code example:

from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras("keras-io/vit-small-ds")

Error

/usr/local/lib/python3.7/dist-packages/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name) 560 if cls is None: 561 raise ValueError( –> 562 f’Unknown {printable_module_name}: {class_name}. Please ensure this ’ 563 'object is passed to the custom_objects argument. See ’ 564 ‘https://www.tensorflow.org/guide/keras/save_and_serialize’

ValueError: Unknown optimizer: Addons>AdamW. Please ensure this object is passed to the custom_objects argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

[gante]

I agree with what @ariG23498 wrote about custom layers, its flexibility (which makes it hard to automate) can be a boon. And the creators of custom layers are mostly power user anyways, in my experience 😄

For completeness of discussion, there is a point yet to be addressed in this discussion. The error that @osanseviero originally points at can be avoided by importing tfa, which is needed to load the optimizer. In other words:

• this works

import tensorflow_addons as tfa
from huggingface_hub import from_pretrained_keras
loaded_model = from_pretrained_keras("keras-io/vit-small-ds")
loaded_model.summary()

• this doesn’t work

from huggingface_hub import from_pretrained_keras
loaded_model = from_pretrained_keras("keras-io/vit-small-ds")
loaded_model.summary()

Digging deeper, we can see that push_to_hub_keras uses tf.keras.models.save_model under the hood (here). We might want to set its include_optimizer argument to False, which removes the optimizer object before serialization, preventing errors like this from optimizers that are not in the standard tensorflow library.

What do you think?

EDIT: at the very least, we can throw a warning when the user is pushing a model to the hub with these kind of optimizers.

[ariG23498]

Hey @osanseviero @merveenoyan

I was successful in uploading the custom objects with push_to_hub_keras API. The main steps are:

• Have a get_config method for all the custom layers.
• Serialization of tensors that are used in the custom layers.

class MultiHeadAttentionLSA(tf.keras.layers.MultiHeadAttention):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # The trainable temperature term. The initial value is
        # the square root of the key dimension.
        self.tau = tf.Variable(
            math.sqrt(float(self._key_dim)),
            trainable=True
        )
        # Build the diagonal attention mask
        diag_attn_mask = 1 - tf.eye(NUM_PATCHES)
        self.diag_attn_mask = tf.cast([diag_attn_mask], dtype=tf.int8)
    
    def get_config(self):
        config = super().get_config()
        config.update({
            "tau": self.tau.numpy(),                       #<---- IMPORTANT
            "diag_attn_mask": self.diag_attn_mask.numpy(), #<---- IMPORTANT
        })
        return config

• Not using augmentation layers inside the model. TensorFlow 2.7 has a problem with serializing the augmentation layers. Here I have used the map function to map the tf.data and used the augmentation as a preprocessing step rather than inside the mode.

Colab Notebook used

Colab Notebook

Usage of the pre-trained vit-ds-small model

loaded_model = from_pretrained_keras("keras-io/vit-small-ds")

_, accuracy, top_5_accuracy = loaded_model.evaluate(test_ds)
print(f"Test accuracy: {round(accuracy * 100, 2)}%")
print(f"Test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%")

Note: You have to use the test augmentation to preprocess the input images before sending it to the model.