添加新Token后加载TensorFlow保存的BERT模型出现形状不匹配错误

错误原因 #

当你给BERT分词器添加新Token并调整词嵌入大小后，模型的词嵌入权重形状已从原始的(30523, 768)变为(30524, 768)（对应新增1个Token）。但直接使用tf.keras.models.load_model()加载.h5模型时，会默认按照原始BERT的结构初始化模型，导致试图将新形状的权重分配给旧形状的变量，触发形状不匹配错误。

解决方案 #

方法1：使用Transformers原生保存/加载方式（推荐） #

Transformers库提供了专门的模型保存方法，能自动处理词嵌入层的调整，避免形状问题：

• 保存模型和分词器：

# 保存微调后的模型
model.save_pretrained("./my_finetuned_bert")
# 同步保存分词器（包含新增的Token）
tokenizer.save_pretrained("./my_finetuned_bert")

• 加载模型和分词器：

from transformers import TFBertModel, BertTokenizer

# 加载保存的分词器
loaded_tokenizer = BertTokenizer.from_pretrained("./my_finetuned_bert")
# 加载保存的模型，自动适配调整后的词嵌入层
loaded_model = TFBertModel.from_pretrained("./my_finetuned_bert")

方法2：手动重建模型结构后加载权重（兼容.h5格式） #

如果必须使用.h5格式保存，需要先重建与微调时完全一致的模型结构，再加载权重：

from transformers import TFBertModel, BertTokenizer

# 1. 重建模型结构，和微调前步骤完全一致
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
# 添加同样的新Token
tokenizer.add_tokens(['new_token'])
# 加载原始预训练模型
bert = TFBertModel.from_pretrained("bert-base-uncased")
# 调整词嵌入大小，匹配分词器的新词汇量
bert.resize_token_embeddings(len(tokenizer))

# 2. 加载.h5格式的权重
bert.load_weights("path_to_model.h5")

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