Pytorch中的交叉熵损失函数nn.CrossEntropyLoss之中包含两个非常有用的参数:
- weight:用来平衡样本之间的不平衡;
- reduction:用来指定损失结果返回的是mean、sum还是none。
事实的使用方法是:
- 使用reduction='mean’的同时指定weights的时候,不会产生任何错误,而且reduction='mean’也是默认的参数值;
- 使用reduction='none’的同时指定weights的时候,会导致产生的结果有错误,这是因为当我们指定reduction='none’的时候,该函数不会对结果使用weight进行规范化。
解决方法如下:
criterion = nn.CrossEntropyLoss(weight=weights, reduction='none')
loss = criterion(x, y)
loss = loss.sum() / weights[y].sum()
这样一来,我们便使用weight进行了规范化处理,从而可以得到正确的结果。
我们可以通过以下代码验证:
x = torch.randn(10, 5)
y= torch.randint(0, 5, (10,))
weights = torch.tensor([1., 2., 3., 4., 5.])
criterion_good = nn.CrossEntropyLoss(weight=weights)
loss_good = criterion_good(x, y)
criterion = nn.CrossEntropyLoss(weight=weights, reduction='none')
loss = criterion(x, y)
loss = loss.sum() / weights[y].sum()
print(loss == loss_good)