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

本文主要是介绍深度学习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 = ksuper(CustomLayer, self).__init__(**kwargs)def get_config(self):config = super().get_config()config["k"] = self.kreturn configdef call(self, input):return tf.multiply(input, 2)

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

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