深度学习图像分割入门：从零理解CLDice损失函数及其PyTorch实现

1次阅读

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

在医学图像分割任务中（比如肿瘤区域划分），我们最常用的 Dice 损失函数存在两个明显痛点：

对小目标不敏感 ：当分割目标只占图像的很小部分时（如早期病灶），Dice 系数会因背景像素主导而失去判别力
拓扑结构破坏 ：传统 Dice 只关注像素级重叠度，可能导致预测结果出现空洞或断裂（如下图肝脏血管该连通的地方断开）

# 传统 Dice 系数计算公式
dice = 2 * |X ∩ Y| / (|X| + |Y|)  # X 为预测值，Y 为真实标签

CLDice 在 Dice 基础上引入连通性惩罚项，其数学表达式分为两部分：

传统 Dice 项 ：保持区域重叠精度
连通性惩罚项 ：强制预测结果与真实标签具有相似拓扑结构

完整公式：

CLDice = α*Dice + (1-α)*[∑(c∈C) Dice(connected(c,X), connected(c,Y)) ]

其中：
– connected(c,X) 表示对预测结果 X 中第 c 个连通分量的提取
– α 是平衡权重（通常取 0.5）

以下是经过显存优化的向量化实现（支持 batch 处理）：

import torch
import torch.nn as nn
import scipy.ndimage as ndimage

def find_connected_components(mask):
    """GPU 加速的连通域查找"""
    # 此处使用 scipy 的 CPU 实现作为示例，实际可用 cc3d 等 GPU 库
    return torch.from_numpy(ndimage.label(mask.cpu().numpy())[0])

class CLDiceLoss(nn.Module):
    def __init__(self, alpha=0.5):
        super().__init__()
        self.alpha = alpha
        self.smooth = 1e-6

    def forward(self, pred, target):
        # 二值化处理
        pred_bin = (pred > 0.5).float()
        target_bin = (target > 0.5).float()

        # 计算传统 Dice
        intersection = (pred_bin * target_bin).sum()
        dice_coef = (2.*intersection + self.smooth) / 
                   (pred_bin.sum() + target_bin.sum() + self.smooth)

        # 连通性计算
        pred_cc = find_connected_components(pred_bin)
        target_cc = find_connected_components(target_bin)

        # 连通性 Dice（简化版实现）conn_dice = 0
        for label in torch.unique(target_cc):
            if label == 0: continue  # 跳过背景
            conn_mask = (pred_cc == label).float()
            conn_target = (target_cc == label).float()
            conn_intersect = (conn_mask * conn_target).sum()
            conn_dice += 2*conn_intersect / (conn_mask.sum() + conn_target.sum() + self.smooth)

        # 组合损失
        loss = self.alpha*(1-dice_coef) + (1-self.alpha)*(1-conn_dice)
        return loss

关键优化点说明：