如何在PyTorch的CrossEntropyLoss中为批量样本设置不同权重？

我知道torch.nn.CrossEntropyLoss会为批量里的每个样本计算损失值，再求和得到最终损失。现在想在求和前，把批量中3个样本的损失分别和数组multiplier里的对应元素相乘，该怎么实现？

注意：这个函数的weight参数是对不同目标类别加权，不是我要的效果。

示例代码：

import torch
import numpy as np
multiplier = torch.from_numpy(np.array([1.0, -2.0, 5.0]))
source = torch.from_numpy(
    np.array([[0.5, -0.6],
              [-3.0, -2.0],
              [-4.0, 2.3]]))
target = torch.from_numpy(np.array([0, 1, 0]))
loss_fn = torch.nn.CrossEntropyLoss()
loss_fn(source, target)

解决方案 #

通过设置CrossEntropyLoss的reduction='none'，可以获取每个样本单独的损失值，之后再和multiplier逐元素相乘并求和即可：

import torch
import numpy as np

multiplier = torch.from_numpy(np.array([1.0, -2.0, 5.0]))
source = torch.from_numpy(
    np.array([[0.5, -0.6],
              [-3.0, -2.0],
              [-4.0, 2.3]]))
target = torch.from_numpy(np.array([0, 1, 0]))

# 配置损失函数为返回每个样本的损失，不自动求和/平均
loss_fn = torch.nn.CrossEntropyLoss(reduction='none')
per_sample_loss = loss_fn(source, target)
# 对每个样本损失加权后求和
weighted_loss = (per_sample_loss * multiplier).sum()

print(weighted_loss)

说明 #

• reduction='none'：让损失函数输出与批量大小一致的张量，每个元素对应单个样本的损失，而非直接输出求和或平均后的总损失。
• 逐元素乘法后调用.sum()，就能得到每个样本损失按对应权重加权后的总和，完全符合需求。

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