Keras自定义模型保存与加载的注意事项

在使用Keras构建深度学习模型时，自定义模型(Model子类)是一种常见的做法。然而，在保存和加载自定义模型时，开发者可能会遇到一些棘手的问题。本文将深入探讨这些问题及其解决方案。

问题背景

当开发者尝试保存并重新加载一个自定义的Keras模型时，可能会遇到以下错误信息：

1. Functional.__init__() got multiple values for keyword argument 'inputs'
2. Function.__init__() got an unexpected keyword argument 'layers'

这些错误通常出现在Keras 3.5.0及更早版本中，主要原因是模型序列化和反序列化过程中的实现差异。

自定义模型实现分析

让我们先看一个典型的问题案例。开发者创建了一个继承自keras.Model的DummyModel，它包含多个卷积层和池化层：

@keras.saving.register_keras_serializable()
class DummyModel(keras.Model):
    def __init__(self, *, input_shape=(28,28,1), filters=[16,32], activation='relu', **kwargs):
        input_spec = keras.layers.Input(shape=input_shape)
        x = input_spec
        # 构建模型结构
        x = layers.Conv2D(filters[0], 3, activation=activation)(x)
        x = layers.Conv2D(filters[1], 3, activation=activation)(x)
        x = layers.MaxPooling2D(3)(x)
        x = layers.Conv2D(filters[1], 3, activation=activation)(x)
        x = layers.Conv2D(filters[0], 3, activation=activation)(x)
        x = layers.GlobalMaxPooling2D()(x)
        super().__init__(inputs=input_spec, outputs=x, **kwargs)
        
        self.filters = filters
        self.activation = activation

    def get_config(self):
        config = super().get_config()
        config.update({
            "input_shape": self.input_shape[1:],
            "filters": self.filters,
            "activation": self.activation,
        })
        return config

问题根源

在Keras 3.5.0及更早版本中，get_config()方法会返回整个计算图的表示，包括所有层的配置信息。这导致在模型反序列化时出现冲突：

1. 当调用super().get_config()时，父类会返回完整的模型结构配置
2. 开发者又添加了自定义配置项，导致配置信息冗余
3. 在加载模型时，Keras无法正确处理这种混合配置

解决方案

对于Keras 3.5.0及更早版本

在旧版本中，需要完全重写get_config()方法，避免调用父类方法：

def get_config(self):
    return {
        "name": self.name,
        "input_shape": self.input_shape[1:],
        "filters": self.filters,
        "activation": self.activation,
    }

同时，必须实现from_config类方法：

@classmethod
def from_config(cls, config):
    return cls(**config)

对于Keras 3.8.0及更新版本

新版本已经修复了这个问题，开发者可以安全地调用super().get_config()并添加自定义配置项，如最初的代码示例所示。

最佳实践

无论使用哪个版本的Keras，以下实践都能确保自定义模型的正确序列化：

1. 始终使用@keras.saving.register_keras_serializable()装饰器注册自定义模型
2. 明确定义get_config()方法，返回所有必要的配置信息
3. 实现from_config类方法，确保能正确重建模型
4. 在保存模型前，先测试模型的保存和加载流程
5. 保持Keras版本的一致性，避免在不同环境中使用不同版本

总结

Keras自定义模型的保存和加载是一个需要特别注意的功能。理解Keras内部如何处理模型配置对于构建可靠的自定义模型至关重要。随着Keras的发展，这些问题正在逐步解决，但了解其背后的机制仍然有助于开发者更好地使用这个强大的深度学习框架。

对于生产环境，建议使用较新的Keras版本(3.8.0+)，以获得更好的开发体验和更稳定的功能支持。如果必须使用旧版本，则需要按照本文介绍的解决方案进行调整。