带权重交叉熵损失手动计算与PyTorch官方实现结果不一致问题求助

核心错误点 #

你手动实现的带权重交叉熵存在两处逻辑错误：

1. 权重匹配逻辑错误：没有为每个样本匹配其对应标签类别的权重，错误地对所有输出维度的损失都乘了权重总和
2. 平均规则错误：PyTorch官方F.cross_entropy设置reduction='mean'时，带权重的平均分母是所有参与计算样本的对应权重之和，不是所有类别的权重总和，也不是样本总数

官方带权重交叉熵计算逻辑 #

• 第一步：计算每个样本的无权重交叉熵损失 loss_i = -log_softmax(pred_i)[label_i]
• 第二步：为每个样本损失乘对应标签类别的权重 w_i = weights[label_i]
• 第三步：加权求和后除以权重总和 total_loss = sum(w_i * loss_i) / sum(w_i)

正确实现代码 #

import torch
import torch.nn.functional as F

# 原始变量定义
pred = torch.tensor([[8,5,3,2,6,1,6,8,4],[2,5,1,3,4,6,2,2,6],[1,1,5,8,9,2,5,2,8],[2,2,6,4,1,1,7,8,3],[2,2,2,7,1,7,3,4,9]]).float()
label = torch.tensor([[3],[7],[8],[2],[5]],dtype=torch.int64)
weights = torch.tensor([1,1,1,10,1,6,1,1,1],dtype=torch.float32)

# 手动实现带权重交叉熵
one_hot = torch.zeros_like(pred).scatter(1, label.view(-1, 1), 1)
log_prb = F.log_softmax(pred, dim=1)
# 计算单样本无权重损失
per_sample_loss = -(one_hot * log_prb).sum(dim=1)
# 匹配每个样本对应的权重
sample_weights = weights[label.view(-1)]
# 加权平均计算最终损失
weighted_loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum()

print(weighted_loss) # 输出为4.7894，和官方结果完全一致

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