微调GPT-2时遇stack期望张量尺寸一致错误的解决方法

问题背景 #

我尝试用个人信息微调GPT-2模型，搭建一个可以让他人了解我的聊天机器人，但运行代码时遇到以下错误：

RuntimeError: stack expects each tensor to be equal size, but got [47] at entry 0 and [36] at entry 1

错误原因是输入文本编码后的张量长度不一致，我的两个样本输入如下：

What is the webisite of ABC company ? -> https://abcdef.org/
Do you know the website of ABC company ? -> It is https://abcdef.org/

问题分析 #

原代码存在两个核心问题：

1. 在Dataset的__getitem__方法中使用tokenizer.encode时，padding=True仅会对单条文本做自适应填充，不会统一到max_length，导致不同样本的张量长度不一致，DataLoader堆叠时触发尺寸不匹配错误。
2. 虽然手动添加了[PAD] token，但没有更新GPT-2模型的embedding层，模型无法识别新增的pad token，后续训练会出现异常。

修复方案及代码 #

以下是修正后的完整代码，解决了上述问题：

import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from torch.utils.data import Dataset, DataLoader

class QADataset(Dataset):
    def __init__(self, questions, answers, tokenizer, max_length):
        self.questions = questions
        self.answers = answers
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.questions)

    def __getitem__(self, index):
        question = self.questions[index]
        answer = self.answers[index]
        input_text = f"Q: {question} A: {answer}"
        
        # 使用tokenizer的调用方法，统一将文本编码到指定max_length
        encoding = self.tokenizer(
            input_text,
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt"
        )
        
        # 移除batch维度，返回单条数据的张量
        input_ids = encoding["input_ids"].squeeze()
        attention_mask = encoding["attention_mask"].squeeze()
        
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask
        }

# 初始化tokenizer，用GPT2的eos_token作为pad token（原生无pad token）
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token

# 加载模型并更新embedding层，适配新增的pad token
model = GPT2LMHeadModel.from_pretrained('gpt2')
model.resize_token_embeddings(len(tokenizer))

# 加载问答数据
questions = ["What is the webisite of ABC company ?", "Do you know the website of ABC company ?"]
answers = ["https://abcdef.org/", "It is https://abcdef.org/"]

# 创建数据集和数据加载器，样本量小，batch_size设为2
max_length = 64
dataset = QADataset(questions, answers, tokenizer, max_length=max_length)
data_loader = DataLoader(dataset, batch_size=2, shuffle=True)

# 训练配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

# 开始训练
for epoch in range(3):
    running_loss = 0.0
    model.train()
    for batch in data_loader:
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        
        # 前向传播，传入attention_mask避免模型关注pad token
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=input_ids
        )
        
        loss = outputs.loss
        
        # 反向传播与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
    
    avg_loss = running_loss / len(data_loader)
    print(f"Epoch {epoch + 1} 平均损失: {avg_loss:.4f}")

# 保存微调后的模型和tokenizer
model.save_pretrained("qa_finetuned_gpt2")
tokenizer.save_pretrained("qa_finetuned_gpt2")

关键修改说明 #

• 替换tokenizer.encode为tokenizer()调用：通过padding="max_length"强制所有样本编码后长度一致，从根源解决张量尺寸不匹配问题。
• 适配pad token：GPT2原生没有pad token，使用eos_token替代，并调用model.resize_token_embeddings更新模型embedding层，确保模型能正确处理pad token。
• 添加attention_mask：训练时传入该参数，让模型忽略pad token位置的计算，提升训练效果。
• 调整batch_size：样本量仅为2，将batch_size设为2避免DataLoader报错。

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