TensorFlow2.0保存模型参数时报错的解决办法

TensorFlow2.0保存模型参数时报错的解决办法

NotImplementedError: Layers with arguments in __init__ must override get_config.

在Tensorflow2.0保存模型时，通常用model.save语句保存模型

model.save("./saver/model_mnist.h5")

有时会出现下面这样的错误，报错信息很长只截取了上下部分

Traceback (most recent call last):
  File "F:/Python_ChengXu/Pycharm Projects/Python_TF2_projects/code4-6.py", line 56, in <module>
    model.save("./saver/the_model_mnist_2.h5")
  File "D:\Program Files (x86)\Anaconda3-5.2.0\envs\tf2\lib\site-packages\tensorflow_core\python\keras\engine\network.py", line 975, in save
    signatures, options)
  File "D:\Program Files (x86)\Anaconda3-5.2.0\envs\tf2\lib\site-packages\tensorflow_core\python\keras\engine\base_layer.py", line 573, in get_config
    raise NotImplementedError('Layers with arguments in `__init__` must '
NotImplementedError: Layers with arguments in `__init__` must override `get_config`.

但有时候却可以正常保存。
究其原因，是因为如果构建的网络中有的层自己定义的网络层时，那在保存模型时就要格外注意。当然，如果神经网络各个层都是用函数直接用的层，就比如下面这样，这样就可以直接用model.save()函数保存，参数是要保存的路径，这样就是以.h5的格式保存，不仅保存了网络的模型，还有训练之后的参数。

input_xs = tf.keras.Input([28,28,1])
conv = tf.keras.layers.Conv2D(32,3,padding = "SAME",activation = tf.nn.relu)(input_xs)
conv = tf.keras.layers.BatchNormalization()(conv) 
conv = tf.keras.layers.Conv2D(64,3,padding = "SAME",activation = tf.nn.relu)(conv)
conv = tf.keras.layers.MaxPool2D(strides = [1,1])(conv)
conv = tf.keras.layers.Conv2D(128,3,padding = "SAME",activation = tf.nn.relu)(conv)

flat = tf.keras.layers.Flatten()(conv)
dense = tf.keras.layers.Dense(512,activation = tf.nn.relu)(flat)
logits = tf.keras.layers.Dense(10,activation = tf.nn.softmax)(dense)

model = tf.keras.Model(inputs = input_xs,outputs = logits)

但是，如果里面有自己定义的网络层，就像这样

class MyLayer(tf.keras.layers.Layer):
    def __init__(self,kernel_size,filter):  #设定层中参数，目前需要定义的参数是卷积核大小和卷积核数目
        self.filter = filter
        self.kernel_size = kernel_size
        super(MyLayer,self).__init__()
    def build(self,input_shape):
        self.weight = tf.Variable(tf.random.normal([self.kernel_size,self.kernel_size,input_shape[-1],self.filter]),name = 'Weight')
        self.bias = tf.Variable(tf.random.normal([self.filter]),name = 'Bias')
        super(MyLayer,self).build(input_shape)
    def call(self,input_tensor):
        conv = tf.nn.conv2d(input_tensor,self.weight,strides = [1,2,2,1],padding = "SAME")
        conv = tf.nn.bias_add(conv,self.bias)
        out = tf.nn.relu(conv) + conv
        return out

input_xs = tf.keras.Input([28,28,1])
conv = tf.keras.layers.Conv2D(32,3,padding = "SAME",activation = tf.nn.relu)(input_xs)
conv = MyLayer(32,3)(conv)
conv = tf.keras.layers.BatchNormalization()(conv)

显然网络层中的MyLayer(32,3)就是自己定义的，这时候在模型保存时会出现上面提到的错误。需要做两点

• 重新get_config()函数。在原来的网络层定义的__init__，build，call函数下面重写get_config()函数，就像这样

    def get_config(self):  #在有自定义网络层时，需要保存模型时，重写get_config函数
        config = {"kernel_size": self.kernel_size,"filter":self.filter}
        base_config = super(MyLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

• 自定义层中的Variable类型的参数需要命名。这个不必多说，build函数里的Variable类型的变量直接name一下就可以，注意是build函数所有Variable类型的变量都要命名，下面只是截取的其中一个

self.bias = tf.Variable(tf.random.normal([self.filter]),name = 'Bias')

完成这两步，再次运行就可以成功保存模型啦。下面运行之后的提示信息

118/118 [==============================] - 94s 799ms/step - loss: 0.3067 - accuracy: 0.9097
模型保存成功
last scaore: [0.10443486766889691, 0.9664]
totally cost 106.19307065010071

我爱小倩！