如何在PyTorch中实现适配不同重要性训练样本的加权损失函数

有非常简便的实现方式，核心思路是先不做损失聚合，得到每个样本的独立损失值，和对应样本权重相乘后再按需聚合即可，具体实现如下：

方法1：直接使用内置MSELoss调整reduction参数 #

PyTorch的nn.MSELoss支持通过reduction参数控制损失输出形式，设置为none时会返回和输入shape一致的每个样本的单独损失，直接乘权重即可：

import torch
import torch.nn as nn

# 定义损失时关闭默认聚合
mse = nn.MSELoss(reduction="none")

# 模拟输入：真实标签、模型预测、样本权重，三者shape完全一致
y_true = torch.tensor([2.0, 5.0, 7.0])
y_pred = torch.tensor([1.8, 5.3, 6.9])
sample_weights = torch.tensor([1.0, 3.0, 0.5]) # 第二个样本优先级最高，错误惩罚最大

# 计算加权损失
per_sample_loss = mse(y_pred, y_true)
weighted_loss = (per_sample_loss * sample_weights).mean() # 也可以用.sum()根据需求聚合

# 后续正常执行反向传播、优化器更新即可
weighted_loss.backward()

方法2：封装自定义加权损失类 #

如果需要频繁使用该加权逻辑，可以直接封装为独立的损失类：

class WeightedMSELoss(nn.Module):
    def __init__(self, reduction: str = "mean"):
        super().__init__()
        self.reduction = reduction

    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
        # 逐样本计算损失并加权
        per_sample_loss = (y_pred - y_true) ** 2
        weighted_loss = per_sample_loss * weights
        # 按配置聚合损失
        if self.reduction == "mean":
            return weighted_loss.mean()
        elif self.reduction == "sum":
            return weighted_loss.sum()
        else:
            return weighted_loss

使用时直接调用自定义损失类即可：

weighted_mse = WeightedMSELoss(reduction="mean")
loss = weighted_mse(y_pred, y_true, sample_weights)
loss.backward()

注意事项 #

• 样本权重的shape需要和单样本损失的shape完全匹配，批量输入、多维度输出场景下也适用，无需额外调整
• 本方案仅针对样本级加权，如果是类别级加权需求，部分损失函数（如nn.CrossEntropyLoss）内置了weight参数可直接传入类权重
• 加权后的损失反向传播、优化器更新逻辑和普通训练完全一致，不需要额外修改

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