Saving Keras models with Custom Layers

Question

I am trying to save a Keras model in a H5 file. The Keras model has a custom layer. When I try to restore the model, I get the following er

Answer 1

You can provide manually the mapping custom_objects in the load_model method as mentioned in the answer https://stackoverflow.com/a/62326857/8056572 but it can be tedious when you have a lot of custom layers (or any custom callables defined. e.g. metrics, losses, optimizers, ...).

Tensorflow provides a utils function to do it automatically: tf.keras.utils.register_keras_serializable

You have to update your CustomLayer as follows:

import tensorflow as tf

@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, **kwargs):
        self.k = k
        super(CustomLayer, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["k"] = self.k
        return config

    def call(self, input):
        return tf.multiply(input, 2)

---

Here is the complete working code:

import tensorflow as tf


@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, **kwargs):
        self.k = k
        super(CustomLayer, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["k"] = self.k
        return config

    def call(self, input):
        return tf.multiply(input, 2)


def main():
    model = tf.keras.models.Sequential(
        [
            tf.keras.Input(name='input_layer', shape=(10,)),
            CustomLayer(10, name='custom_layer'),
            tf.keras.layers.Dense(1, activation='sigmoid', name='output_layer')
        ]
    )
    print("SUMMARY OF THE MODEL CREATED")
    print("-" * 60)
    print(model.summary())
    model.save('model.h5')

    del model

    print()
    print()

    model = tf.keras.models.load_model('model.h5')
    print("SUMMARY OF THE MODEL LOADED")
    print("-" * 60)
    print(model.summary())

if __name__ == "__main__":
    main()

And the corresponding output:

SUMMARY OF THE MODEL CREATED
------------------------------------------------------------
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
custom_layer (CustomLayer)   (None, 10)                0         
_________________________________________________________________
output_layer (Dense)         (None, 1)                 11        
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
_________________________________________________________________
None


WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
SUMMARY OF THE MODEL LOADED
------------------------------------------------------------
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
custom_layer (CustomLayer)   (None, 10)                0         
_________________________________________________________________
output_layer (Dense)         (None, 1)                 11        
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
_________________________________________________________________
None