PyTorch目标检测：不平衡数据集的带权重训练方案咨询

你当前使用的TorchVision Faster RCNN默认损失是在模型的roi_heads模块中计算的，分类损失采用不带权重的交叉熵损失。针对不平衡数据集，我们可以通过给分类损失添加类别权重来优化，具体步骤如下：

1. 计算类别权重 #

先统计训练数据集中每个类别的目标数量，再基于逆频率（或平方逆频率）计算对应权重，代码示例：

import torch

def calculate_class_weights(train_dataset, num_classes, ignore_background=True):
    # 初始化类别计数数组
    class_counts = torch.zeros(num_classes, dtype=torch.float32)
    # 遍历数据集统计每个类别的目标数量
    for idx in range(len(train_dataset)):
        _, targets = train_dataset[idx]
        labels = targets['labels']
        for label in labels:
            class_counts[label.item()] += 1
    # 避免类别计数为0导致除以0，给空类别设为1
    class_counts[class_counts == 0] = 1
    total_samples = class_counts.sum()
    
    if ignore_background:
        # 背景类（类别0）权重设为1，仅给自定义类别计算权重
        weights = torch.ones(num_classes, dtype=torch.float32)
        # 从类别1开始计算逆频率权重
        for cls in range(1, num_classes):
            weights[cls] = total_samples / ((num_classes - 1) * class_counts[cls])
    else:
        # 所有类别（包括背景）计算逆频率权重
        weights = total_samples / (num_classes * class_counts)
    
    return weights.to(DEVICE)

# 计算类别权重
class_weights = calculate_class_weights(train_dataset, num_classes, ignore_background=True)

2. 修改模型的分类损失函数 #

TorchVision的Faster RCNN默认分类损失不带权重，我们通过猴子补丁替换为带权重的交叉熵损失：

import torch.nn.functional as F
from torchvision.models.detection.roi_heads import fastrcnn_loss

# 创建带权重的交叉熵损失函数
cls_loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights)

def weighted_fastrcnn_loss(class_logits, box_regression, labels, regression_targets):
    labels = torch.cat(labels, dim=0)
    regression_targets = torch.cat(regression_targets, dim=0)
    
    # 使用带权重的分类损失
    classification_loss = cls_loss_fn(class_logits, labels)
    
    # 边界框回归损失保持默认逻辑不变
    sampled_pos_inds_subset = torch.where(labels > 0)[0]
    labels_pos = labels[sampled_pos_inds_subset]
    N, num_classes = class_logits.shape
    box_regression = box_regression.reshape(N, -1, 4)

    box_loss = F.smooth_l1_loss(
        box_regression[sampled_pos_inds_subset, labels_pos],
        regression_targets[sampled_pos_inds_subset],
        beta=1 / 9,
        reduction="sum",
    )
    box_loss = box_loss / labels.numel()

    return classification_loss, box_loss

# 替换模型原有的损失计算函数
model.roi_heads.fastrcnn_loss = weighted_fastrcnn_loss

3. 保持原有训练流程不变 #

你的训练代码无需额外修改，模型在训练时会自动使用带权重的分类损失计算loss_dict中的loss_classifier项，后续的损失求和、反向传播逻辑完全兼容。

补充说明 #

• 如果数据集的边界框数量也存在严重不平衡（比如某些类别的框极少），可以考虑给边界框回归损失也添加权重，但通常分类损失的权重调整对不平衡数据集的性能提升更明显。
• 权重计算方式可按需调整，比如使用平方逆频率（weights = torch.sqrt(total_samples / ((num_classes - 1) * class_counts[cls]))），避免少数类权重过大导致模型过拟合。

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