如何在PyTorch中修改预训练模型的自注意力层数与多头注意力头数

Hugging Face Transformers库中所有预训练模型的结构参数都保存在对应Config对象中，你只需要先修改配置参数，再加载模型即可，具体操作如下：

1 加载并修改模型配置 #

首先加载roberta-base的默认配置，再修改两个核心参数：

• num_hidden_layers：对应Transformer编码器的总层数，也就是你所说的自注意力层数，roberta-base默认值为12
• num_attention_heads：对应每一层多头注意力的头数，roberta-base默认值为12

将你原有代码中加载模型的部分替换为如下代码即可：

from transformers import AutoModelForSequenceClassification, AutoConfig

# 加载roberta-base原始配置
config = AutoConfig.from_pretrained(PRETRAINED_MODEL_NAME)
# 自定义自注意力层数，示例改为8层
config.num_hidden_layers = 8
# 自定义多头注意力头数，示例改为8头
# 注意：头数必须能整除隐藏层维度，roberta-base隐藏层维度为768，因此头数只能取768的正约数
config.num_attention_heads = 8

# 基于修改后的配置加载模型，和原结构匹配的参数会继承预训练权重，改动部分随机初始化
model = AutoModelForSequenceClassification.from_pretrained(PRETRAINED_MODEL_NAME, config=config).to(device)

2 注意事项 #

• 修改结构后如果仍加载预训练权重，只有和原始roberta-base结构匹配的参数会继承预训练值，改动的部分会随机初始化，需要在你的讽刺分类数据集上充分训练才能收敛
• 如果不需要继承预训练权重，想从零训练自定义结构的模型，可以直接使用model = AutoModelForSequenceClassification(config).to(device)加载，无需传入预训练模型名称
• 除了上述两个参数外，你也可以通过修改config的其他字段自定义模型结构，比如hidden_size修改隐藏层维度、hidden_dropout_prob修改dropout概率等

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