BERT预训练模型新增token调用resize_token_embeddings报错咨询

问题背景 #

• 持有已训练完成的模型权重：models/BERT-pretrain-1-step-5000.pkl
• 需求：给分词器新增特殊Token[TRA]，调整模型嵌入层尺寸匹配新词表
• 运行调整代码时报错：AttributeError: 'BertModel' object has no attribute 'resize_token_embeddings'
• 原错误推测：自定义预训练权重的嵌入层尺寸不匹配

报错根因 #

之前的尺寸不匹配推测不成立，核心原因是导入BertModel使用的第三方旧包pytorch_pretrained_bert_inset的模型实现中，根本没有内置resize_token_embeddings方法，和加载的自定义权重没有关系。

可落地实现方案 #

方案1：统一依赖库版本（推荐，稳定性最高） #

不要混用不同来源的模型和分词器实现，全部统一用transformers库的官方实现，自带成熟的嵌入层调整接口，代码如下：

import torch
# 统一从transformers导入模型和分词器，弃用旧的第三方pytorch_pretrained_bert相关包
from transformers import AutoTokenizer, BertModel
from torch.nn.utils.rnn import pad_sequence
import tqdm

# 第一步：加载分词器，新增特殊Token
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
num_added_toks = tokenizer.add_tokens(['[TRA]'], special_tokens=True)
tra_token_id = tokenizer.convert_tokens_to_ids('[TRA]')
print(f'[TRA] token id: {tra_token_id}')

# 第二步：初始化基础模型，先调整嵌入层尺寸匹配新词表
model_bert = BertModel.from_pretrained('bert-base-uncased')
model_bert.resize_token_embeddings(len(tokenizer))

# 第三步：加载自定义训练权重，自动兼容尺寸差异
custom_weight = torch.load('models/BERT-pretrain-1-step-5000.pkl', map_location=torch.device('cpu'))
model_weight = model_bert.state_dict()
# 过滤掉尺寸不匹配的旧嵌入层权重，其余权重正常加载
filtered_weight = {k:v for k,v in custom_weight.items() if k in model_weight and v.shape == model_weight[k].shape}
model_weight.update(filtered_weight)
model_bert.load_state_dict(model_weight)

# 第四步：初始化新增Token的嵌入（用已有词嵌入均值初始化，收敛更快）
with torch.no_grad():
    input_emb = model_bert.get_input_embeddings()
    input_emb.weight[tra_token_id] = torch.mean(input_emb.weight[:-num_added_toks], dim=0)

方案2：保留原有旧版BertModel依赖 #

如果因为项目兼容要求不能替换模型导入源，不需要调用不存在的内置方法，手动实现嵌入层替换即可：

import torch
from pytorch_pretrained_bert_inset import BertModel
from transformers import AutoTokenizer

# 加载分词器、新增特殊Token
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
num_added_toks = tokenizer.add_tokens(['[TRA]'], special_tokens=True)
new_vocab_size = len(tokenizer)
tra_token_id = tokenizer.convert_tokens_to_ids('[TRA]')

# 加载已训练的自定义模型
model_bert = BertModel.from_pretrained(
    'bert-base-uncased', 
    state_dict=torch.load('models/BERT-pretrain-1-step-5000.pkl', map_location=torch.device('cpu'))
)
old_emb_layer = model_bert.embeddings.word_embeddings
old_vocab_size, embed_dim = old_emb_layer.weight.shape

# 手动创建匹配新词表尺寸的嵌入层
new_emb_layer = torch.nn.Embedding(
    num_embeddings=new_vocab_size,
    embedding_dim=embed_dim,
    padding_idx=old_emb_layer.padding_idx
)
# 复制原有词嵌入权重
new_emb_layer.weight.data[:old_vocab_size] = old_emb_layer.weight.data
# 初始化新增Token的嵌入
new_emb_layer.weight.data[tra_token_id] = torch.mean(old_emb_layer.weight.data, dim=0)
# 替换模型原有嵌入层
model_bert.embeddings.word_embeddings = new_emb_layer

注意事项 #

• 不要跨库混用模型和分词器实现，极易出现方法缺失、权重键名不匹配、尺寸对齐错误等问题
• 调整嵌入层尺寸的操作要放在加载自定义权重之前，减少权重兼容处理的复杂度
• 新增特殊Token不要使用默认随机初始化，采用已有词嵌入的均值初始化可以大幅降低后续微调的成本

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