Keras中EfficientNet模型保存与加载的性能问题分析

问题背景

在使用Keras 3和TensorFlow 2.17版本时，开发者发现了一个关于EfficientNet模型保存与加载的性能问题。当使用.keras格式保存一个可训练的EfficientNet基础模型后，重新加载模型时会出现性能显著下降的情况，而使用.h5格式则不会出现此问题。

问题重现

通过简化测试代码可以清晰地重现这个问题：

import numpy as np
import keras

# 构建基于EfficientNetB0的自定义模型
inputs = keras.Input((224, 224, 3))
base_model = keras.applications.EfficientNetB0(weights="imagenet", include_top=False)
x = base_model(inputs)
outputs = keras.layers.Dense(10, activation="softmax")(x)
model = keras.Model(inputs, outputs)

# 保存并重新加载模型
model.save("model.keras")
revived_model = keras.models.load_model("model.keras")

# 测试预测一致性
images = np.random.uniform(size=(2, 224, 224, 3)).astype("float32")
outputs = model.predict(images)
revived_outputs = revived_model.predict(images)
np.testing.assert_allclose(outputs, revived_outputs)  # 此处会抛出异常

根本原因分析

经过深入调查，发现问题出在keras.layers.Normalization层的状态处理上。这个层有两类重要参数：

1. 权重参数：包括adapted_mean和adapted_std，这些会被正常保存和加载
2. 状态变量：包括mean和std，这些是实际用于预测的变量，但不会被自动保存为权重

关键问题在于：

• 当使用.h5格式保存时，会触发finalize_state()方法，正确更新状态变量
• 但使用新的.keras格式保存时，特别是加载ImageNet预训练权重时，finalize_state()不会被自动调用

解决方案

目前有两种可行的解决方案：

1. 临时解决方案：在加载模型后手动调用finalize_state()

base_model.layers[2].finalize_state()

2. 长期解决方案：等待Keras团队修复这个问题，确保在使用.keras格式时也能正确处理Normalization层的状态

影响范围

这个问题主要影响：

• 使用EfficientNet系列模型（B0-B7）
• 使用Keras 3和TensorFlow 2.17及以上版本
• 使用.keras格式保存模型
• 加载了ImageNet预训练权重的情况

值得注意的是，EfficientNetV2和MobileNet系列模型不受此问题影响。

最佳实践建议

对于当前版本，建议开发者：

1. 如果必须使用EfficientNet，考虑暂时使用.h5格式保存模型
2. 或者在使用.keras格式后，手动调用finalize_state()
3. 关注Keras官方更新，及时获取修复版本

这个问题展示了深度学习框架中模型序列化与反序列化的复杂性，特别是在处理有状态的层时。理解这些底层机制有助于开发者更好地调试和解决类似问题。