如何pickle序列化tensorflow.keras自定义评估指标？

• Keras 自定义指标不支持原生pickle序列化，你遇到的AttributeError: 'MulticlassMeanIoU' object has no attribute 'update_state_fn'报错，本质是Keras指标在初始化阶段会动态生成图执行专用的状态更新函数，这类运行时生成的TensorFlow函数、设备绑定张量属性，不在pickle的序列化支持范围内，和你自定义类的写法没有直接关系。
• 不需要强行用pickle保存指标，Keras本身提供了成熟的自定义对象序列化方案，稳定性远高于pickle，还能实现和模型的绑定存储、跨会话加载。

这种方式不需要单独维护指标文件，模型加载后可直接还原绑定的指标实例，不会出现训练端和加载端代码逻辑不一致的问题。

1. 给自定义类加上Keras序列化注册装饰器，修正冗余参数

   import tensorflow as tf
   N_CLASSES = 15
   
   # 注册为Keras可序列化自定义对象
   @tf.keras.utils.register_keras_serializable(package="CustomMetrics")
   class MulticlassMeanIoU(tf.keras.metrics.MeanIoU):
       def __init__(self,
                   num_classes = None,
                   name = "Multi_MeanIoU",
                   dtype = None):
           # 移除__init__中无实际作用的y_true、y_pred占位参数
           super().__init__(num_classes=num_classes, name=name, dtype=dtype)
   
       def get_config(self):
           # 父类方法已经会自动收集num_classes、name、dtype等初始化参数，直接返回即可
           return super().get_config()
   
       def update_state(self, y_true, y_pred, sample_weight = None):
           y_pred = tf.math.argmax(y_pred, axis=-1)
           return super().update_state(y_true, y_pred, sample_weight)
   

2. 模型编译时直接传入自定义指标实例，训练完成后保存为Keras原生格式

   # 模型编译
   model.compile(
       optimizer="adam",
       loss="sparse_categorical_crossentropy",
       metrics=[MulticlassMeanIoU(num_classes=N_CLASSES)]
   )
   # 模型训练逻辑省略...
   
   # 保存为TF2.10+推荐的.keras格式，会自动存储绑定的指标配置
   model.save("/path/to/your/trained_model.keras")
   

3. 跨会话加载模型时，只要提前导入自定义类，即可自动还原所有绑定的指标，不需要额外传参

   # 从你维护的公共自定义模块导入指标类，保证所有脚本用同一份定义
   from my_project.custom_metrics import MulticlassMeanIoU
   
   # 加载模型，自动还原网络结构、权重、优化器状态、绑定的指标
   loaded_model = tf.keras.models.load_model("/path/to/your/trained_model.keras")
   # 直接获取还原后的指标实例即可使用
   restored_metric = loaded_model.metrics[0]
   

注意：把自定义指标类存到项目公共的工具模块中，所有训练、推理脚本统一从这个模块导入，就不会出现多份代码逻辑不同步的问题。

如果需要单独保存、加载指标实例，不要用pickle，用Keras原生的配置+权重存储方案：

1. 保存指标

   import json
   met = MulticlassMeanIoU(num_classes=N_CLASSES)
   
   # 保存可JSON序列化的指标初始化配置
   with open("/path/to/metric_config.json", "w") as f:
       json.dump(met.get_config(), f)
   # 保存指标的状态权重
   met.save_weights("/path/to/metric_weights.h5")
   

2. 加载指标

   import json
   from my_project.custom_metrics import MulticlassMeanIoU
   
   # 从配置初始化实例
   with open("/path/to/metric_config.json", "r") as f:
       met_config = json.load(f)
   restored_met = MulticlassMeanIoU.from_config(met_config)
   # 加载权重还原指标状态
   restored_met.load_weights("/path/to/metric_weights.h5")
   

就算通过重写类的__getstate__、__setstate__方法手动补全update_state_fn这类缺失属性，pickle序列化的产物也会强依赖保存时的TensorFlow版本、设备环境，跨版本、跨GPU/CPU环境加载时极易出现兼容问题，没有生产可用性。

内容的提问来源于stack exchange，提问作者Manuel Popp