使用PyTorch训练Transformer时批量输入的损失计算报错排查

核心问题 #

你遇到的两个错误根源在于两点：一是模型输出维度不符合序列语言模型的要求，二是对CrossEntropyLoss的输入格式及padding处理逻辑理解有误。

分步解决 #

1. 修正模型输出维度 #

Transformer语言模型需要为输入序列中的每个token输出对应位置的预测结果，因此模型的最终输出维度必须是 [batch_size, seq_len, num_classes]（其中seq_len是当前batch的最长序列长度，num_classes是词汇表大小）。

检查你的Transformer实现，确保最后一层线性层的输出维度正确——它需要将模型的隐藏状态（维度通常为[batch_size, seq_len, hidden_dim]）映射到词汇表大小，输出[batch_size, seq_len, num_classes]。

2. 正确计算带padding过滤的损失 #

当模型输出维度正确后，按照以下步骤修改训练循环，即可解决损失计算的问题：

# Training loop
model.train()  # Set the model to training mode
for epoch in range(num_epochs):
    running_loss = 0.0
    for batch in dataloader:
        inputs = batch.to(device)
        optimizer.zero_grad()

        # Forward pass - 模型输出维度应为 [batch_size, seq_len, num_classes]
        outputs = model(inputs)

        # 将输出和标签展平为一维，对应每个token的预测与真实标签
        outputs_flat = outputs.view(-1, num_classes)  # shape: [batch_size * seq_len, num_classes]
        targets_flat = inputs.view(-1)                # shape: [batch_size * seq_len]

        # 创建mask过滤padding的0，只保留有效token的位置
        valid_mask = (targets_flat != 0)
        # 仅对有效token计算损失
        loss = criterion(outputs_flat[valid_mask], targets_flat[valid_mask])

        # 反向传播与优化
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(dataloader)
    print(f"Epoch {epoch+1}/{num_epochs} Loss: {epoch_loss:.4f}")

print("Training done :)")

3. 错误原因解析 #

• 第一个错误（RuntimeError）：你的模型输出是[32, 10007]（每个样本仅对应一个预测结果），但输入标签是[32, seq_len]（每个样本包含多个token标签）。CrossEntropyLoss默认要求每个输入样本对应一个标签（1D target），不支持“多标签”输入，因此触发报错。
• 第二个错误（ValueError）：你传入损失函数的outputs是整个batch的输出（batch_size=32），但masked_inputs.unsqueeze(0)是单个样本的标签（batch_size=1），两者batch维度不匹配，导致损失函数报错。

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