BCE和Dice混合损失函数实战：解决图像分割中的类别不平衡问题

1次阅读

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

在图像分割任务中，我们经常会遇到类别不平衡的问题。具体表现为某些类别的像素数量远远多于其他类别，比如医学图像中的病灶区域通常只占整张图像的很小一部分。这种不平衡会导致模型倾向于预测占比较大的类别，从而忽略小目标的识别。

传统二元交叉熵损失（BCE, Binary Cross Entropy）在这种情况下表现不佳，因为它平等对待所有像素的预测误差。而 Dice 系数（Dice Coefficient）则通过对预测和真实标签的交集与并集的比值来衡量相似度，天然倾向于关注重叠区域，因此能更好地处理类别不平衡问题。

BCE 损失的数学表达式为：

$$
L_{BCE} = -\frac{1}{N}\sum_{i=1}^N [y_i\log(p_i) + (1-y_i)\log(1-p_i)]
$$

优点：
– 计算简单
– 梯度稳定

缺点：
– 对类别不平衡敏感
– 小目标预测效果差

Dice 系数的定义为：

$$
Dice = \frac{2|X \cap Y|}{|X| + |Y|}
$$

对应的 Dice 损失为：

$$
L_{Dice} = 1 – Dice
$$

优点：
– 对类别不平衡鲁棒
– 直接优化分割区域的重叠

缺点：
– 训练初期可能不稳定
– 对小误差敏感

结合 BCE 和 Dice 的混合损失：

$$
L_{Mix} = \alpha L_{BCE} + \beta L_{Dice}
$$

其中 α 和 β 是可调节的权重参数。这种组合可以同时利用 BCE 的稳定性和 Dice 对类别不平衡的鲁棒性。

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

class BCEDiceLoss(nn.Module):
    """
    BCE 和 Dice 混合损失函数
    Args:
        alpha (float): BCE 损失的权重
        beta (float): Dice 损失的权重
        smooth (float): 防止除零的平滑项
        reduction (str): 损失聚合方式，'mean' 或 'sum'
    """def __init__(self, alpha=0.5, beta=0.5, smooth=1e-6, reduction='mean'):
        super(BCEDiceLoss, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.smooth = smooth
        self.reduction = reduction
        self.bce = nn.BCEWithLogitsLoss(reduction=reduction)

    def forward(self, inputs, targets):
        """
        Args:
            inputs (torch.Tensor): 模型输出，未经 sigmoid
            targets (torch.Tensor): 真实标签，值在 0 - 1 之间
        """
        # BCE 损失
        bce_loss = self.bce(inputs, targets)

        # Dice 损失
        probs = torch.sigmoid(inputs)
        intersection = (probs * targets).sum()
        union = probs.sum() + targets.sum()

        # 添加平滑项防止除零
        dice_loss = 1 - (2. * intersection + self.smooth) / (union + self.smooth)

        # 混合损失
        total_loss = self.alpha * bce_loss + self.beta * dice_loss

        return total_loss

我们在一个医学图像分割数据集上进行了实验，比较了三种损失函数的表现：