深度学习Keras保存模型（当包含自定义层时）

20210409 -

0.引言

一般来说，如果没有什么特殊情况，那么在进行保存模型的时候，通过调用一些api进行保存即可，在文章《深度学习的基础知识与问题汇总》简单介绍了一种方式，直接保存模型，并重新载入。

但是在今天的实验中，出现的需求就是，如果定义的模型中，包含了用户自定义的层就会报错，具体情况见[1]，在保存模型的时候没有问题，但是载入时就会报错。
简单说明一下实验环境

python 3.6.8
tensorflow-gpu 2.3.1
Keras 2.4.3

1. 加载模型报错：未定义层

在[1]中，如果没有对自定义的层进行一系列的规定，那么在加载模型的时候， 就会报错为定义层。

ValueError: Unknown layer: CustomLayer

上面这种问题使用的api是load_model，通过这种方式加载整个模型以及各种权值，针对这种错误，可以通过两种方法来解决。在问答[2]中都提到了，一种是在加载模型的时候，在api指定自定义类，如下：

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

可能如果你使用了自定义的损失函数，也需要将这部分内容传输进去。
而另外一种方法更方便，针对如果定义了多个自定义内容：

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)

在自定义层上加上修饰器。

注意看，这里他自定义类的时候，上面这些要实现的函数。最重要的是，将自己的参数在函数get_config中进行保存。具体可以看[2]的方式，可以以他的编程方式作为模板，将模型保存起来，其中还包含了参数初始化的内容，而且可以看到在build过程中引入的add_weight不用在get_config中声明。

2. 保存权值

本次实验中，最后使用的方式是仅仅保存权值，反正逻辑上都得先定义这个模型。使用的方式是在训练的时候加入了保存模型的回调函数。

参考

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