如何在PyTorch CrossEntropyLoss中为多类别多输出任务设置专属类别权重

原生PyTorch的CrossEntropyLoss无法直接实现为每个输出列单独设置类别权重——它的weight参数是全局生效的，会对所有输出分支的同类别的样本应用相同权重。不过你可以通过以下两种方式实现需求：

方法1：拆分损失计算，为每个输出分支单独初始化损失函数 #

针对每个输出列，分别创建一个CrossEntropyLoss实例，传入该列对应的类别权重，最后将各分支的损失求和（或加权求和）。

代码示例 #

import torch
import torch.nn as nn

# 1. 预计算各输出列的类别权重（根据你的数据分布）
output1_weights = torch.tensor([0.425/0.6, 0.15/0.1, 0.425/0.3], dtype=torch.float32)
output2_weights = torch.tensor([0.425/0.4, 0.15/0.3, 0.425/0.3], dtype=torch.float32)
output3_weights = torch.tensor([# 填入OUTPUT_3对应的类别权重], dtype=torch.float32)

# 2. 为每个输出分支定义独立的损失函数
loss_fn1 = nn.CrossEntropyLoss(weight=output1_weights)
loss_fn2 = nn.CrossEntropyLoss(weight=output2_weights)
loss_fn3 = nn.CrossEntropyLoss(weight=output3_weights)

# 3. 前向传播与损失计算
# 假设模型输出三个分支的logits，每个形状为(batch_size, 3)
logits1, logits2, logits3 = model(inputs)
# 假设标签对应三个分支，每个形状为(batch_size,)
target1, target2, target3 = targets

# 计算各分支损失并求和
loss1 = loss_fn1(logits1, target1)
loss2 = loss_fn2(logits2, target2)
loss3 = loss_fn3(logits3, target3)
total_loss = loss1 + loss2 + loss3  # 可根据任务需求调整各分支损失的权重占比

# 反向传播
total_loss.backward()

方法2：自定义多输出交叉熵损失函数 #

通过自定义nn.Module，统一处理所有输出分支的加权损失计算，代码更简洁封装。

代码示例 #

import torch
import torch.nn as nn

class MultiOutputCrossEntropyLoss(nn.Module):
    def __init__(self, branch_weights):
        super().__init__()
        self.branch_weights = branch_weights  # 列表，每个元素对应一个输出列的类别权重张量
        self.base_loss = nn.CrossEntropyLoss(reduction='none')
    
    def forward(self, logits, targets):
        # logits: (batch_size, num_outputs, num_classes)
        # targets: (batch_size, num_outputs)
        total_loss = 0.0
        num_outputs = logits.shape[1]
        
        for idx in range(num_outputs):
            # 提取当前分支的logits和标签
            branch_logits = logits[:, idx, :]
            branch_target = targets[:, idx]
            # 计算无归约的基础损失
            raw_loss = self.base_loss(branch_logits, branch_target)
            # 获取当前分支中每个样本对应类别的权重
            class_weights = self.branch_weights[idx][branch_target]
            # 计算加权后的分支损失并累加
            weighted_loss = (raw_loss * class_weights).mean()
            total_loss += weighted_loss
        
        return total_loss

# 使用自定义损失函数
branch_weights = [output1_weights, output2_weights, output3_weights]
loss_fn = MultiOutputCrossEntropyLoss(branch_weights)

# 假设模型输出形状为(batch_size, 3, 3)，标签形状为(batch_size, 3)
logits = model(inputs)
targets = ...
total_loss = loss_fn(logits, targets)
total_loss.backward()

两种方法都能精准实现为每个输出列的每个类别单独设置权重的需求，可根据你的代码结构选择合适的方式。

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