GPT2训练报错TypeError：仅单元素整数张量可转为索引的解决方法

问题代码 #

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

train_data_path = 'train_dataset.txt'

with open(train_data_path, 'r', encoding='utf-8') as file:
    train_text = file.read()

train_encodings = tokenizer(train_text, return_tensors='pt', truncation=True, max_length=128)

train_dataset = TextDataset(
    tokenizer=tokenizer,
    file_path=train_encodings['input_ids'],
    block_size=128,
)

train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

for epoch in range(3):  # Adjust the number of epochs as needed
    for batch in train_loader:
        inputs = batch["input_ids"]
        labels = batch["labels"]
        outputs = model(inputs, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

model.save_pretrained('D:/Data Science/2. Notebooks/')
tokenizer.save_pretrained('D:/Data Science/2. Notebooks/')

报错位置及信息 #

报错行：

train_dataset = TextDataset(tokenizer=tokenizer, file_path=train_encodings['input_ids'], block_size=128,)       
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)

错误信息：

TypeError: only integer tensors of a single element can be converted to an index

train_encodings内容：

{'input_ids': tensor([[   42,    72, 16175,  1134, 10662,  1878,    80,  1031,   264, 30102,
             430, 12379, 33133, 21681, 30102,    77,   324,   133,   247,    74,
             334, 15201, 30102,    81,   198,  2202, 21681, 30102,    77,   331,
             263,   521,   133,   247, 28808,    72,   819,    75,   133,   247,
              81,   256,  1134,   346,    76,   133,   247,    75,   312,   343,
             198,   128,   108,   358,    72,   277,  1134, 12769,   299,   133,
             247, 15908,   390,    88, 17062,   133,   247,    77,   198,    48,
           30102,    89, 30102,    75,   479, 25236,   264,   361,   133,   247,
           28781,   133,   247,    81,   220,   133,   247, 18809,   133,   247,
           1225,   133,   247,    81,   133,   247,    74, 28808,    72,   512,
           21681,   331, 34174,    76, 30102, 46481, 21681,   198, 34458, 46481,
           20402,   465,    82,   133,   247,    75,   133,   247,  4372,   133,
             247,  1312, 46481,    75,   133,   247,    81,   308]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1]])}

错误原因 #

TextDataset的file_path参数要求传入原始文本文件的路径字符串，它会自动读取文件内容并完成tokenize和分块操作。但你传入的是已经tokenize后的tensor（train_encodings['input_ids']），参数类型不匹配，导致内部处理索引时抛出类型错误。

解决方案 #

方案1：直接使用TextDataset加载原始文本文件（推荐） #

这是TextDataset的标准用法，无需手动提前tokenize，它会自动处理所有流程：

# 替换原train_dataset创建代码
train_dataset = TextDataset(
    tokenizer=tokenizer,
    file_path=train_data_path,  # 传入原始文本文件路径
    block_size=128,
)

train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)

后续训练循环无需修改，保持原代码即可。

方案2：用TensorDataset处理已编码的数据 #

如果你已经提前完成了tokenize操作，可以用TensorDataset替代TextDataset，但需要手动处理维度和标签：

from torch.utils.data import TensorDataset

# 提取input_ids并调整维度（原input_ids是[1, 128]，去掉多余的batch维度）
input_ids = train_encodings['input_ids'].squeeze(0)
# LM任务中标签与输入一致
labels = input_ids.clone()

# 创建TensorDataset
train_dataset = TensorDataset(input_ids, labels)
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)

# 调整训练循环，确保输入维度符合模型要求（模型接受[batch_size, seq_len]格式）
for epoch in range(3):
    for batch in train_loader:
        inputs, labels = batch
        # 增加batch维度
        inputs = inputs.unsqueeze(0)
        labels = labels.unsqueeze(0)
        outputs = model(inputs, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

注意事项 #

• 若训练数据量较大，优先选择方案1，TextDataset会自动处理文本分块、批量tokenize，避免手动处理维度问题。
• 方案2仅适合小批量已编码数据，需严格确保输入维度与模型要求匹配。

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