加载训练好的BERT模型时遇TFBertModel未知层错误求助

问题原因 #

你的模型中使用了transformers库提供的TFBertModel自定义层，用Keras的model.save()保存后，加载时Keras无法自动识别该外部层，必须通过custom_objects参数显式传入该层的定义。你当前的加载代码错误地传入了自定义模型函数，而非实际需要的TFBertModel。

解决方案 #

1. 修正模型加载代码
   确保加载前已导入TFBertModel和Adam（模型编译时用到了该优化器），然后修改load_model的custom_objects参数：

# 确保导入必要模块
from transformers import TFBertModel
from tensorflow.keras.optimizers import Adam

# 修正后的加载代码
loaded_model = tf.keras.models.load_model(
    'bert_model.h5', 
    custom_objects={"TFBertModel": TFBertModel, "Adam": Adam}
)

2. 可选：更稳定的保存/加载方式
   如果后续仍遇到自定义层兼容性问题，推荐使用SavedModel格式保存整个模型，这种格式对复杂层结构的支持更好：

# 保存为SavedModel格式（替代h5）
model.save('bert_model_savedmodel')

# 加载SavedModel格式模型
loaded_model = tf.keras.models.load_model(
    'bert_model_savedmodel',
    custom_objects={"TFBertModel": TFBertModel, "Adam": Adam}
)

验证 #

加载完成后，可通过loaded_model.summary()确认模型结构是否正常，再进行预测测试。

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