微调google/reformer-enwik8做NER时axial_pos_shape参数报错求解

错误根本原因 #

你两次尝试失败的核心原因是修改了和预训练权重强绑定的参数：hidden_size和axial_pos_embds_dim。原始预训练模型的词嵌入、轴向位置嵌入、Transformer层、输出层的权重维度都和hidden_size=1024、axial_pos_embds_dim=(256,768)（256+768=1024）严格对应，修改这两个参数必然导致权重维度不匹配。
另外你选择的ReformerModelWithLMHead是带语言模型生成头的模型，天生适配文本生成任务，做NER这种token分类任务应该用ReformerForTokenClassification，避免输出头维度不匹配的问题。 #

可行解决方法 #

方法1：仅调整轴向位置形状（保留预训练权重） #

不需要修改hidden_size和axial_pos_embds_dim，仅通过切片位置嵌入权重、调整axial_pos_shape的方式缩小支持的最大序列长度，不会破坏预训练权重的匹配性：

import transformers
import torch

# 加载模型时仅指定num_labels，不要修改其他结构参数
model = transformers.ReformerForTokenClassification.from_pretrained(
    "google/reformer-enwik8", 
    num_labels=9
)

# 自定义目标最大序列长度，比如你可以根据数据集的最大长度设为2048，拆分为两个因数相乘
target_max_len = 2048
target_axial_shape = (16, 128) # 16*128=2048，可根据你的需要调整为任意两个乘积等于target_max_len的正整数

# 切片位置嵌入权重，仅保留对应长度的参数，维度保持和预训练一致
model.reformer.embeddings.position_embeddings.weights[0] = torch.nn.Parameter(
    model.reformer.embeddings.position_embeddings.weights[0][:target_axial_shape[0]]
)
model.reformer.embeddings.position_embeddings.weights[1] = torch.nn.Parameter(
    model.reformer.embeddings.position_embeddings.weights[1][:target_axial_shape[1]]
)

# 仅更新配置中和序列长度相关的参数
model.config.axial_pos_shape = target_axial_shape
model.config.max_position_embeddings = target_max_len

改完后你只需要将输入序列填充到target_max_len即可，内存占用会远低于填充到65536的情况。

方法2：动态适配每个batch的序列长度 #

如果你的数据集样本长度差异很大，可以每个batch单独计算当前batch的最大长度，向上取最近的可拆分为两个正整数乘积的数值，再按方法1的逻辑动态调整轴向形状和位置嵌入权重，进一步降低内存占用。 #

代码优化建议 #

1. 不要调用模型的内部方法_pad_to_mult_of_chunk_length，直接用对应tokenizer的padding参数，将序列填充到你设定的max_position_embeddings即可。
2. 训练时传入attention_mask，模型会自动忽略填充位置的计算，不会影响效果。

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