PyTorch 中交叉熵损失的输入维度

Input dimension for CrossEntropy Loss in PyTorch

对于 batch_size = 1 的二元分类问题,我有 logit 和标签值,我需要使用它们来计算损失。

logit: tensor([0.1198, 0.1911], device='cuda:0', grad_fn=<AddBackward0>)
label: tensor(1], device='cuda:0')
# calculate loss
loss_criterion = nn.CrossEntropyLoss()
loss_criterion.cuda()
loss = loss_criterion( b_logits, b_labels )

然而,这总是导致以下错误,

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

CrossEntropyLoss 实际要求的输入维度是多少?

您传递的张量形状错误。
shape 应该是 (from doc)

  • Input: (N,C) where C = number of classes
  • Target: (N) where each value is 0 ≤ targets[i] ≤ C−1

所以在这里,b_logits 形状应该是 ([1,2]) 而不是 ([2]) 以使其成为正确的形状你可以像 b_logits.view(1,-1) 一样使用 torch.view

b_labels形状应该是([1]).
例如:

b_logits = torch.tensor([0.1198, 0.1911], requires_grad=True)
b_labels = torch.tensor([1])
loss_criterion = nn.CrossEntropyLoss()

loss = loss_criterion( b_logits.view(1,-1), b_labels )
loss
tensor(0.6581, grad_fn=<NllLossBackward>)