20210409 -

0. introduction

Generally speaking , If there are no special circumstances , When saving the model , By calling some api Just save it , In the article 《 Summary of basic knowledge and problems of in-depth learning 》 A way is briefly introduced , Save the model directly , And reload .

But in today's experiment , The emerging demand is , If in the defined model , If the user-defined layer is included, an error will be reported , See [1], There is no problem saving the model , But an error will be reported when loading .
Briefly explain the experimental environment

python 3.6.8
tensorflow-gpu 2.3.1
Keras 2.4.3

1. Error in loading model ： Undefined layer

stay [1] in , If there is no series of regulations for the customized layer , When loading the model , An error will be reported as the definition layer .

ValueError: Unknown layer: CustomLayer

The above question uses api yes load_model, Load the whole model and various weights in this way , For this kind of mistake , There are two ways to solve . Questions and answers [2] It's all mentioned in , One is when loading the model , stay api Specify custom classes , as follows ：

new_model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})

Maybe if you use a custom loss function , You also need to transfer this part of the content .
The other method is more convenient , If multiple customized contents are defined ：

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)

Add modifiers to the custom layer .

Look for the , Here, when he customizes the class , The above functions to be implemented . most important of all , Put your own parameters in the function get_config To save . For details, see [2] The way , You can use his programming method as a template , Save the model , It also contains the content of parameter initialization , And you can see in build Introduced in the process add_weight No use get_config In a statement .

2. Save weights

In this experiment , The last way to use is to just save the weights , In any case, we have to define this model first . The way to use is to add a callback function to save the model during training .

Reference resources

[1]Saving Keras models with Custom Layers
[2]Not able to load a saved model with custom layer