Keras模型序列化与自定义损失函数加载问题解析

问题背景

在Keras深度学习框架中，从3.6版本升级到3.7版本后，用户报告了一个关于模型加载的问题：当尝试加载在Keras 3.6中保存的包含自定义损失函数的模型时，系统会抛出"Could not locate function"错误。这个问题特别出现在使用@tf.keras.utils.register_keras_serializable装饰器注册的自定义损失函数上。

技术细节分析

序列化机制的变化

在Keras 3.6版本中，自定义损失函数的序列化方式存在一个关键特性：它仅使用函数名称作为键值进行序列化，而忽略了装饰器中指定的包名。这种实现方式虽然简化了序列化过程，但可能导致命名冲突和加载问题。

Keras 3.7版本通过提交795df4ef63566ec869fe2512373f4346e1e02746修复了这个问题，使序列化过程能够正确处理包名信息。然而，这一改进也移除了向后兼容的备用代码路径，导致在3.6版本中保存的模型无法在3.7及以上版本中正确加载。

损失函数与激活函数的差异

值得注意的是，这个问题仅影响自定义损失函数，而不影响自定义激活函数。这是因为Keras中对这两类函数的序列化处理方式存在差异：

1. 损失函数序列化：在3.6版本中仅使用函数名作为键
2. 激活函数序列化：始终使用get_registered_name方法生成包含包名的完整键

这种不一致的处理方式是导致版本间兼容性问题的根本原因。

解决方案

临时解决方案

对于需要从Keras 3.6加载模型的用户，可以采用以下方法：

# 重新定义自定义损失函数
@keras.saving.register_keras_serializable("attention_loss")
def attention_loss(y_actual, y_predicted):
    # 实现代码

# 加载模型时显式指定custom_objects
reloaded_model = keras.models.load_model(
    "custom_model.keras",
    custom_objects={"attention_loss": attention_loss},
)

这种方法通过显式提供自定义对象映射，绕过了自动反序列化机制的问题。

长期建议

1. 统一升级环境：尽可能在开发和部署环境中使用相同的Keras版本
2. 完整保存自定义对象：考虑将自定义对象的实现代码与模型一起保存
3. 使用标准组件：优先使用Keras内置的损失函数和激活函数

技术启示

这个问题揭示了深度学习框架版本管理中的几个重要考量：

1. 序列化兼容性：框架升级时应考虑模型文件的向后兼容性
2. API一致性：相似功能的组件(如损失函数和激活函数)应保持一致的序列化策略
3. 文档完整性：变更可能影响用户工作流的特性时，应在文档中明确说明

对于Keras用户而言，理解模型的序列化机制和版本间差异对于长期维护模型部署至关重要。在升级框架版本时，建议进行充分的测试，特别是对于包含自定义组件的模型。