TensorFlow Keras SavedModel 在保存和加载两次后抛出 TypeError

Question

当我创建具有一个或多个自定义层的 Keras 模型时，我可以使用 model.save() 方法以使用 TensorFlow SavedModel 格式保存 Keras 模型。

我可以使用tf.keras.models.load_model() 函数从文件系统中加载这个模型并再次将其保存到文件系统中。

但是当我第二次从文件系统加载 SavedModel 时，它会失败并出现以下异常：

TypeError: f(inputs, training, training, training, training, *, training, training) missing 1 required argument: training

您可以尝试使用以下代码复制此问题：

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):
    def call(self, inputs, *args, **kwargs):
        return inputs

model1 = tf.keras.Sequential([
    CustomLayer()
])
model1.build((None, 1))
model1.compile()
model1.save("model1")

model2 = tf.keras.models.load_model("model1")
model2.save("model2")

# This line should raise a TypeError.
model3 = tf.keras.models.load_model("model2")

Answer 1

为什么会出现问题

问题在于 TensorFlow SavedModel 格式实际上并没有序列化自定义 Python 代码。它只保存自定义 Keras 层和其他 Python 对象生成的 TensorFlow 图。

tf.keras.models.load_model() 函数——默认情况下——不返回 Python 层。相反，它返回一个占位符层，其中包含 TensorFlow 计算图的相同部分。我们可以在我的问题的示例中看到这一点：

>>> model1.layers
[<__main__.CustomLayer at 0x7ff04c14ee20>]

>>> model2.layers
[<keras.saving.saved_model.load.CustomLayer at 0x7ff114fd7be0>]

当model2 被保存并从文件系统加载时，TensorFlow 无法正确解析CustomLayer.call() 中的*args 和**kwargs 参数。

我不知道实际的错误是在保存代码中，还是在加载代码中，或者两者兼而有之。

真正的修复需要在 TensorFlow/Keras 中进行，但与此同时，也有

解决方法

您可以选择以下任何一种解决方法来避免自定义 Keras 层的序列化错误。

更改Layer.call()上的签名

目前Layer.call()上的官方方法签名是def call(self, inputs, *args, **kwargs):

但 TensorFlow 在尝试加载具有此签名的自定义层的模型时会抛出 TypeError。要修复错误，请使用签名def call(self, inputs): 编写所有自定义层。如果你的层在训练或推理过程中表现不同，那么你可以使用方法签名def call(self, inputs, training=None):

这使得 TensorFlow 可以更轻松地生成在 keras.saving.saved_model.load 模块中生成的占位符层。但是这个占位符层仍然与原始 Python 代码不完全相同。

在tf.keras.models.load_model() 上使用custom_objects 参数

可以使用原始 Python 层而不是占位符层加载模型。只需将字典映射层名称传递给 Python 层类对象。这要求您的代码能够导入原始 Python 层。我的问题中的示例可以固定如下：

model3 = tf.keras.models.load_model(
    "model2",
    custom_objects=dict(
        CustomLayer=CustomLayer,
    ),
)

确保您的图层实现Layer.get_config() 并返回一个字典，其中包含从头开始重新创建图层所需的所有参数。该层必须能够使用Layer.from_config() 重新创建。

导入 Python 层并将其添加到 Keras 的全局注册表中

Keras 维护自定义 Python 类和其他对象的全局注册表，以便在加载 SavedModel 时引用。您可以使用 @tf.keras.utils.register_keras_serializable() 装饰器注册您的自定义 Keras 层。例如：

@tf.keras.utils.register_keras_serializable(
   package="my_python_package"
)
class CustomLayer(tf.keras.layers.Layer):
    def call(self, inputs, *args, **kwargs):
        return inputs

此方法还要求您的层正确实现Layer.get_config()。

用tf.keras.utils.custom_object_scope()安装Python层对象

与上述两种解决方案非常相似，tf.keras.utils.custom_object_scope() 上下文管理器可以指定反序列化时使用哪些自定义层。