在R中保存加载微调BERT模型遇TokenEmbedding报错的解决方法

这个问题我之前帮不少开发者解决过——本质是Keras的HDF5格式加载器不认识BERT模型里的自定义TokenEmbedding层，给你两个靠谱的解决思路：

方法1：加载时手动注册自定义层（适合坚持用HDF5格式的场景） #

如果你一定要用save_model_hdf5保存，加载的时候得明确告诉Keras这个自定义层的存在：

1. 首先确保你已经导入了包含TokenEmbedding层定义的代码或模块（比如你用的transformers库相关部分）
2. 加载模型时通过custom_objects参数映射这个自定义层：

# 先导入相关库，确保TokenEmbedding层的定义被加载
library(transformers)

# 加载模型并指定自定义层
model2 <- load_model_hdf5("model.h5", custom_objects = list(TokenEmbedding = TokenEmbedding))

要是不确定层的准确类名，可以先打印原模型的层结构确认：

print(model$layers)

找到对应层的类名后再填入custom_objects。

方法2：改用SavedModel格式保存（更推荐，无额外配置） #

HDF5格式对复杂自定义模型的支持一直不太友好，TensorFlow官方推荐的SavedModel格式能完整保存模型的所有结构（包括自定义层），加载时不需要额外处理：

保存模型 #

save_model_tf(model, "bert_finetuned_model")

加载模型 #

model2 <- load_model_tf("bert_finetuned_model")

额外优化建议（针对预训练BERT场景） #

如果你是用transformers库做的微调，更建议用库自带的预训练模型保存/加载方法，这是专门为这类模型设计的，稳定性拉满：

# 保存微调后的模型
model$save_pretrained("bert_finetuned")

# 加载模型（根据你的任务替换模型类，比如BertForTokenClassification等）
model2 <- BertForSequenceClassification$from_pretrained("bert_finetuned")

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