无法保存自定义子类模型

Question

受到tf.keras.Model子类化的启发，我创建了自定义模型。
我可以训练它并获得成功的结果，但我无法保存它。
我使用的是python3.6和tensorflow v1.10（或v1.9）。

完整的最小代码示例在此处:

import tensorflow as tf
from tensorflow.keras.datasets import mnist


class Classifier(tf.keras.Model):
    def __init__(self):
        super().__init__(name="custom_model")

        self.batch_norm1 = tf.layers.BatchNormalization()
        self.conv1 = tf.layers.Conv2D(32, (7, 7))
        self.pool1 = tf.layers.MaxPooling2D((2, 2), (2, 2))

        self.batch_norm2 = tf.layers.BatchNormalization()
        self.conv2 = tf.layers.Conv2D(64, (5, 5))
        self.pool2 = tf.layers.MaxPooling2D((2, 2), (2, 2))

    def call(self, inputs, training=None, mask=None):
        x = self.batch_norm1(inputs)
        x = self.conv1(x)
        x = tf.nn.relu(x)
        x = self.pool1(x)

        x = self.batch_norm2(x)
        x = self.conv2(x)
        x = tf.nn.relu(x)
        x = self.pool2(x)

        return x


if __name__ == '__main__':
    (x_train, y_train), (x_test, y_test) = mnist.load_data()

    x_train = x_train.reshape(*x_train.shape, 1)[:1000]
    y_train = y_train.reshape(*y_train.shape, 1)[:1000]

    x_test = x_test.reshape(*x_test.shape, 1)
    y_test = y_test.reshape(*y_test.shape, 1)

    y_train = tf.keras.utils.to_categorical(y_train)
    y_test = tf.keras.utils.to_categorical(y_test)

    model = Classifier()

    inputs = tf.keras.Input((28, 28, 1))

    x = model(inputs)
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(10, activation="sigmoid")(x)

    model = tf.keras.Model(inputs=inputs, outputs=x)
    model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
    model.fit(x_train, y_train, epochs=1, shuffle=True)

    model.save("./my_model")

错误信息：

1000/1000 [==============================] - 1s 1ms/step - loss: 4.6037 - acc: 0.7025
Traceback (most recent call last):
  File "/home/user/Data/test/python/mnist/mnist_run.py", line 62, in <module>
    model.save("./my_model")
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1278, in save
    save_model(self, filepath, overwrite, include_optimizer)
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/saving.py", line 101, in save_model
    'config': model.get_config()
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1049, in get_config
    layer_config = layer.get_config()
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1028, in get_config
    raise NotImplementedError
NotImplementedError

Process finished with exit code 1

我查看了错误代码并发现 get_config 方法检查了 self._is_graph_network

有人处理过这个问题吗？

谢谢！

更新1：
在keras 2.2.2（不是tf.keras）
找到注释（用于模型保存）
文件: keras/engine/network.py
函数: get_config

# 子类网络不可序列化
# （除非子类网络的作者实现了序列化）。

所以，显然这行不通...
我想知道为什么他们没有在文档中指出它（像：“使用无法保存的子类！”）

更新2：
在keras文档中找到：

在子类化模型中，模型的拓扑结构定义为Python代码
（而不是作为一组层的静态图）。这意味着模型的拓扑结构无法被检查或序列化。因此，以下方法和属性对于子类化模型不可用：

model.inputs 和 model.outputs。
model.to_yaml() 和 model.to_json()
model.get_config() 和 model.save()。

因此，没有办法使用子类化保存模型。
只能使用Model.save_weights()

- RedEyed

---

子类化模型无法序列化的原因是Keras需要跟踪每个张量的历史记录，以确定图形的结构。每个张量都应该是tf.keras.layers.Layer的输出。然而，子类化模型在其call方法中包含了像tf.nn.relu这样的原始TensorFlow操作，因此无法进行序列化。 - Jie.Zhou

那么，如果我只使用 tf.keras 呢？ 答案：这是不起作用的。 - RedEyed

如果您真正想要一个子类化模型，我的建议是忘记使用 Model.save，而是使用 Model.save_weights 仅保存模型的权重，然后使用 Model.load_weights 加载权重。否则，如果您仍然想保存整个模型而不仅仅是权重，则必须遵循 Keras 的函数 API 指南。 - Jie.Zhou

哦，这真的很有帮助。Model.save_weights 运行良好。非常感谢你！ - RedEyed

Answer 1

TensorFlow 2.2

感谢 @cal 提醒我，新版本的 TensorFlow 已经支持保存自定义模型！

通过使用 model.save 方法保存整个模型，并使用 load_model 方法恢复先前存储的子类模型。以下代码片段描述了如何实现它们。

class ThreeLayerMLP(keras.Model):

  def __init__(self, name=None):
    super(ThreeLayerMLP, self).__init__(name=name)
    self.dense_1 = layers.Dense(64, activation='relu', name='dense_1')
    self.dense_2 = layers.Dense(64, activation='relu', name='dense_2')
    self.pred_layer = layers.Dense(10, name='predictions')

  def call(self, inputs):
    x = self.dense_1(inputs)
    x = self.dense_2(x)
    return self.pred_layer(x)

def get_model():
  return ThreeLayerMLP(name='3_layer_mlp')

model = get_model()
# Save the model
model.save('path_to_my_model',save_format='tf')

# Recreate the exact same model purely from the file
new_model = keras.models.load_model('path_to_my_model')

参见：使用Keras保存和序列化模型 - 第二部分：保存和加载子类化模型

TensorFlow 2.0

简而言之：

1. 对于自定义的子类化Keras模型，不要使用model.save()；
2. 请改用save_weights()和load_weights()。

---

在TensorFlow团队的帮助下，得出了保存自定义子类化Keras模型的最佳实践是保存其权重，并在需要时重新加载。

无法简单地保存Keras自定义子类化模型的原因是它包含无法安全序列化的自定义代码。但是，如果我们拥有相同的模型结构和自定义代码，则可以毫无问题地保存/加载模型的权重。

由Keras的作者Francois Chollet编写的一篇优秀的教程，介绍了如何在Tensorflow 2.0中保存/加载Sequential/Functional/Keras/Custom Sub-Class Models，可在这里的Colab中查看。在保存子类化模型部分中，它指出：

序列模型和函数式模型是表示层DAG的数据结构。因此，它们可以安全地进行序列化和反序列化。

子类化模型有所不同，它不是数据结构，而是一段代码。模型的架构是通过call方法的主体定义的。这意味着无法安全地序列化模型的架构。要加载模型，您需要访问创建它的代码（即模型子类的代码）。或者，您可以将该代码序列化为字节码（例如通过pickling），但这是不安全且通常不可移植的。

- Huan

---

请注意，将权重和架构分开保存会丢失编译信息（优化器和损失），这些信息是使用model.save()保存的。 - Daniel Braun

根据tensorflow的文档，现在似乎可以保存包含自定义代码的完整模型了。您的答案可能会被编辑 :) - Cal

TensorFlow 2.2的附加说明：在此之前，您必须调用model._set_inputs。 - prouast

@Huan，我使用model.save("NameOfModel", save_format='tf')保存了模型，但是在使用loaded_model = keras.models.load_model('./NameOfModel')加载模型后，出现了ValueError Python inputs incompatible with input_signature: inputs: ( Tensor("IteratorGetNext:0", shape=(None, 2), dtype=int32)) input_signature: ( TensorSpec(shape=(None, 2), dtype=tf.int64, name='input_1'))的错误提示。 - hanzgs

谢谢您的回答。但是这是否也适用于RNN呢？因为在我的使用情况下，一旦我包括了self.cell = layers.SimpleRNNCell(10) self.rnn = layers.RNN(self.cell, return_sequences=True, return_state=False)，我就无法使用save_model。 - Paco Wong

我得到的是：模型<__main__.ThreeLayerMLP object at 0x17cc20dd0>无法保存，因为输入形状不可用或者模型的前向传递未定义。要定义前向传递，请重写Model.call()。要指定输入形状，可以直接调用build(input_shape)，或者使用Model()、Model.fit()或Model.predict()在实际数据上调用模型。如果您有自定义的训练步骤，请确保通过Model.__call__在训练步骤中调用前向传递，即model(inputs)，而不是model.call()。 - undefined

Answer 2

Tensorflow 2.1允许使用SavedModel格式保存子类模型

从我开始使用Tensorflow时，我一直是模型子类的粉丝，我觉得这种构建模型的方式更具有Python风格和协作友好性。但是使用这种方法保存模型总是一个痛点。

最近，我开始更新我的知识，并了解到以下信息，似乎对于Tensorflow 2.1是正确的：

子类模型

我发现这个

第二种方法是使用model.save保存整个模型，然后使用load_model恢复先前存储的子类模型。

这最后将模型、权重和其他内容保存到了SavedModel文件中

最后确认：

保存自定义对象：如果您正在使用SavedModel格式，则可以跳过此部分。 HDF5和SavedModel之间的关键区别在于HDF5使用对象配置保存模型结构，而SavedModel保存执行图。因此，SavedModels能够保存像子类模型和自定义层这样的自定义对象，而无需要求原始代码。 我亲自测试过，效果很好，对于子类化模型，model.save()会生成SavedModel保存。现在不再需要使用model.save_weights()或相关函数，它们现在更适用于特定用例。 这应该是对于所有对模型子类化感兴趣的人来说痛苦路径的结束。

- ElPapi42

Answer 3

根据 1.13预发布补丁说明，此问题将在即将发布的版本中得到解决。 以下是更新内容： Keras和Python API： - 可以通过tf.contrib.saved_model.save_keras_model保存子类化的Keras模型。 编辑： 但是，事实上这个功能并没有像说明文档那样完善。对于v1.13的该函数文档指出: 模型限制： - 顺序模型和函数式模型始终可以被保存。 - 当serving_only=True时，只有子类化模型才能被保存。这是由于当前实现复制模型以导出训练和评估图表。由于无法确定子类化模型的拓扑结构，因此无法克隆子类化模型。未来子类化模型将完全可导出。

- qwitwa

Answer 4

我找到了解决方法。创建一个新的模型，并从保存的.h5模型中加载权重。虽然这种方式不是首选，但它可以在keras 2.2.4和tensorflow 1.12中使用。

class MyModel(keras.Model):  
     def __init__(self, inputs, *args, **kwargs):
          outputs = func(inputs)
     super(MyModel, self).__init__( inputs=inputs, outputs=outputs, *args, **kwargs)

def get_model():  
    return MyModel(inputs, *args, **kwargs)

model = get_model()
model.save(‘file_path.h5’)

model_new = get_model()  
model_new.compile(optimizer=optimizer, loss=loss, metrics=metrics) 
model_new.load_weights(‘file_path.h5’)  
model_new.evaluate(x_test, y_test, **kwargs)

- SUN Qinggang

Answer 5

更新：7月20日

最近我也尝试创建了自己的子类层和模型。编写自己的get_config()函数可能会有一些困难。所以我使用了model.save_weights(path_to_model_weights)和model.load_weights(path_to_model_weights)。当你想要加载权重时，记得用相同的架构创建模型，然后再调用model.load_weights()。详细信息请参考tensorflow的指南。

旧答案（仍然正确）

实际上，tensorflow文档中提到：

为了保存/加载具有自定义定义层或子类模型的模型，您应该重写get_config方法，并可选地重写from_config方法。此外，您还应该注册自定义对象，以便Keras知道它的存在。

例如：

class Linear(keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super(Linear, self).__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

    def get_config(self):
        config = super(Linear, self).get_config()
        config.update({"units": self.units})
        return config


layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)

输出结果为：

{'name': 'linear_8', 'trainable': True, 'dtype': 'float32', 'units': 64}

你可以尝试运行这段简单的代码。例如，在函数"get_config()"中，删除config.update()，看看会发生什么。若想了解更多详细信息，请参阅this和this。这些是Tensorflow网站上有关Keras指南的内容。

- Heran

---

你写的答案是针对自定义层的，不适用于问题中要求的子类模型。 - Saurabh Kumar

Answer 6

在使用tf.saved_model.save之前，请先使用model.predict

- Antoine Liutkus

---

这与所问问题无关。预测不会有任何作用。 - Saurabh Kumar

Answer 7

重新创建模型

keras.models.load_model('path_to_my_model')

对我没用

首先，我们必须从构建的模型中保存权重

model.save_weights('model_weights', save_format='tf')

然后，我们必须为子类模型初始化一个新实例，然后使用一个记录编译和train_on_batch，并加载构建模型的权重。

loaded_model = ThreeLayerMLP(name='3_layer_mlp')
loaded_model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
loaded_model.train_on_batch(x_train[:1], y_train[:1])
loaded_model.load_weights('model_weights')

这在TensorFlow==2.2.0中完美运行。

- hanzgs