BiSeNet V2损失函数优化实战：解决语义分割中的类别不平衡问题

1次阅读

共计 2111 个字符，预计需要花费 6 分钟才能阅读完成。

语义分割任务中的类别不平衡问题普遍存在。以道路场景为例，路面像素通常占总图像面积的 70% 以上，而交通标志、行人等关键目标的像素占比可能不足 1%。这种不平衡会导致两个典型问题：

模型倾向于预测高频类别，低频类别召回率显著下降
标准交叉熵损失函数 $\mathcal{L}_{CE}=-\sum y_i\log(p_i)$ 会被高频类别主导

在 BiSeNet V2 这类实时分割网络中，由于采用轻量级设计，特征表示能力有限，类别不平衡问题会被进一步放大。实验显示，使用标准 CE 损失时，Cityscapes 数据集中 ”traffic sign” 类别的 IoU 往往低于 20%。

Focal Loss：通过 $(1-p_t)^\gamma$ 动态降低易分类样本的权重
$$\mathcal{L}_{FL}=-(1-p_t)^\gamma \log(p_t)$$
适合解决前景 - 背景极度不平衡场景
Dice Loss：直接优化分割任务的 IoU 指标
$$\mathcal{L}_{Dice}=1-\frac{2\sum p_i y_i}{\sum p_i + \sum y_i}$$
对小目标敏感但训练初期可能不稳定
Lovasz-Softmax：基于凸优化的 IoU 替代函数
数学性质优秀但计算复杂度较高

推荐采用加权组合方式：

$$\mathcal{L}{total} = \mathcal{L}} + \lambda_1\mathcal{L{FL} + \lambda_2\mathcal{L}$$

经验参数设置：

$\lambda_1=1.0$ (Focal Loss)
$\lambda_2=0.5$ (Dice Loss)
Focal Loss 的 $\gamma=2.0$

import torch
import torch.nn as nn
import torch.nn.functional as F

class CompoundLoss(nn.Module):
    def __init__(self, gamma=2.0, dice_weight=0.5):
        super().__init__()
        self.gamma = gamma
        self.dice_weight = dice_weight

    def forward(self, pred, target):
        # 交叉熵损失
        ce_loss = F.cross_entropy(pred, target)

        # Focal Loss
        logpt = F.log_softmax(pred, dim=1)
        pt = torch.exp(logpt)
        fl_loss = -((1 - pt) ** self.gamma) * logpt
        fl_loss = fl_loss.mean()

        # Dice Loss (避免除零)
        smooth = 1e-5
        pred_prob = F.softmax(pred, dim=1)
        target_onehot = F.one_hot(target, num_classes=pred.shape[1]).permute(0,3,1,2)
        intersection = (pred_prob * target_onehot).sum(dim=(2,3))
        union = pred_prob.sum(dim=(2,3)) + target_onehot.sum(dim=(2,3))
        dice_loss = 1 - (2. * intersection + smooth) / (union + smooth)
        dice_loss = dice_loss.mean()

        return ce_loss + fl_loss + self.dice_weight * dice_loss

显存优化技巧：