基于BERT构建的Keras模型保存报Can't pickle及IndexError问题求助

1. pickle序列化失败原因 #

TensorFlow/Keras模型（尤其是包含Hugging Face预训练BERT层的模型）内部存在大量计算图依赖、动态生成的运行时方法，不属于pickle支持的普通Python可序列化对象范畴。报错中提到的LayerCall.__call__方法是模型调用时动态生成的实例方法，和序列化时读取的静态定义方法不匹配，因此无法被pickle序列化存储。注：pickle本身就不是TensorFlow官方推荐的模型存储方案。

2. model.save()触发IndexError原因 #

核心问题出在BERT层输出的取值写法：代码中直接通过[1]下标取BERT模型输出元组的pooler_output，Keras在序列化追踪计算图的过程中，无法正确识别这种对自定义层输出元组的下标索引操作，尤其是TensorFlow 2.3~2.6版本存在该场景的已知兼容bug。
额外存在的风险点：输出层使用sigmoid激活后输出已为0~1区间的概率值，但损失函数设置了from_logits=True，参数不匹配也可能间接引发序列化逻辑异常。

• 修正BERT层输出的取值方式，不要使用下标索引，直接调用Hugging Face输出类的官方属性：

  # 替换原来的embeddings = bert_model([INPUT_IDs, INPUT_MASK])[1]
  embeddings = bert_model([INPUT_IDs, INPUT_MASK]).pooler_output
  

• 修正损失函数参数不匹配问题：

  # 替换原来的loss = BinaryCrossentropy(from_logits=True)
  loss = BinaryCrossentropy(from_logits=False)
  

• 优先使用TensorFlow官方的SavedModel格式存储全模型，不要使用h5格式，存储时无需加后缀：

  model.save("bert_classifier_model")
  

• 加载全模型时，通过custom_objects参数传入BERT自定义层声明：

  from transformers import TFBertModel
  loaded_model = tf.keras.models.load_model("bert_classifier_model", custom_objects={"TFBertModel": TFBertModel})
  

• 若仍存在兼容问题，可选择仅存储模型权重，避免计算图序列化的兼容问题：

  # 存储权重
  model.save_weights("bert_classifier_weights.h5")
  
  # 加载权重：先构建完全一致的模型结构，再执行加载
  model.load_weights("bert_classifier_weights.h5")
  

内容的提问来源于stack exchange，提问作者Mohamed Ahmed