无法保存自定义子类模型

灵感来自tf.keras.Model 子类化 https://www.tensorflow.org/guide/keras#model_subclassing我创建了自定义模型。
我可以训练它并获得成功的结果，但是我无法保存它.
我使用 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

我查看了错误行并发现获取配置方法检查self._is_graph_network

有人处理这个问题吗？

Thanks!

更新1：
在 keras 2.2.2 上（不是 tf.keras）
找到评论（用于模型保存）
文件：keras/engine/network.py
功能：获取配置

# 子类网络不可序列化
#（除非序列化是通过
# 子类网络的作者）。

所以，显然这是行不通的...
我想知道，他们为什么不在书中指出这一点文档 https://www.tensorflow.org/guide/keras（例如：“使用子类化而无法保存！”）

更新2：
在发现keras文档 https://keras.io/models/about-keras-models/:

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

模型.输入和模型.输出。
model.to_yaml() 和 model.to_json()
model.get_config() 和 model.save()。

因此，无法通过使用子类化来保存模型。
可以只使用Model.save_weights()

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')

See: 使用 Keras 保存和序列化模型 - 第二部分：保存和加载子类模型 https://www.tensorflow.org/guide/keras/save_and_serialize#part_ii_saving_and_loading_of_subclassed_models

TensorFlow 2.0

TL;DR:

1. 不使用model.save()用于自定义子类 keras 模型；
2. use save_weights() and load_weights()反而。

在 Tensorflow 团队的帮助下，事实证明保存自定义子类 Keras 模型的最佳实践是保存其权重并在需要时加载回来。

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

Keras 的作者 Francois Chollet 编写了一篇很棒的教程，介绍了如何在 Colab 中的 Tensorflow 2.0 中保存/加载顺序/函数/Keras/自定义子类模型，网址为here https://colab.research.google.com/drive/172D4jishSgE3N7AO6U2OKAA_0wNnrMOq#scrollTo=mJqOn0snzCRy. In 保存子类模型部分，它说：

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

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