PyTorch CrossEntropyLoss报错：Target 2越界问题求助

报错信息 #

File "/home/es/anaconda3/envs/pyg-meta/lib/python3.10/site-packages/torch/nn/modules/loss.py", line 1164, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/es/anaconda3/envs/pyg-meta/lib/python3.10/site-packages/torch/nn/functional.py", line 3014, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
IndexError: Target 2 is out of bounds.

运行时张量形状 #

(Pdb) out.shape
torch.Size([64, 2])
(Pdb) Yt_train.shape
torch.Size([64])

核心代码 #

criterion = torch.nn.CrossEntropyLoss()
for data in loader:
    _, out = model(data.x, data.edge_index, data.batch)
    Yt_train = data.y.type(torch.LongTensor)
    loss = criterion(out, Yt_train)

原因与解决方法 #

• 直接原因：模型输出out的第二个维度是2，对应类别索引范围为0和1，但标签张量Yt_train中存在值为2的样本，超出了索引范围，导致报错。

• 具体解决方式：

  1. 确认任务类别数：如果实际是3分类任务（标签包含0、1、2），修改模型最后一层的输出神经元数量为3，让out.shape变为[64, 3]。
  2. 排查数据集标签：如果任务是2分类，检查数据集里的标签数据，找出值为2的异常样本，要么修正标签为0或1，要么直接过滤这些样本。
  3. 设备一致性检查：确保标签张量和模型输出在同一设备上（比如模型用GPU时，标签也要转到GPU），可以修改标签转换代码为：Yt_train = data.y.type(torch.LongTensor).to(model.device)。

内容的提问来源于stack exchange，提问作者S.EB