图像分割的损失函数

本文翻译自losses for image segmentation

基于分布相似度的损失

在此类损失函数中，主要使用信息论中的交叉熵机制来度量模型输出和真实目标之间所包含的信息的相似性。

交叉熵损失

假设$P(Y=0)=p$，$P(Y=1)=1-p$。预测值由logistic/sigmoid函数计算得出：

$P(\hat Y=0)= \frac{1}{1+e^{-z}}$和$P(\hat Y=1)=1- \frac{1}{1+e^{-z}}=1-\hat p$

交叉熵损失函数的定义形式如下：

$$\left(p, \hat{p}\right) = -\left(p \log\left(\hat{p}\right) + (1-p) \log\left(1 - \hat{p}\right)\right)$$

上述交叉熵损失为二维交叉熵损失。

加权交叉熵损失

加权交叉熵损失（weighted corss entropy, WCE）是交叉熵损失的一种变体，其中，所有的正样本都被乘以一个系数以进行加权。该损失函数常用于类别不平衡问题中。例如，当你有一张有10%的黑像素和90%的白像素的图片时，常规的CE效果不会太好。

WCE定义如下式：

$$\text{WCE}\left(p, \hat{p}\right) = -\left(\beta p \log\left(\hat{p}\right) + (1-p) \log\left(1 - \hat{p}\right)\right)$$

减少假负样本的数目，将$\beta$设置为大于1.增加假正样本的数目，将$\beta$设置为小于1。

weighted bce代码实现：

def criterion_pixel(logit_pixel, truth_pixel):
    logit = logit_pixel.view(-1)
    truth = truth_pixel.view(-1)
    assert(logit.shape==truth.shape)

    loss = F.binary_cross_entropy_with_logits(logit, truth, reduction='none')
    if 0:
        loss = loss.mean()
    if 1:
        pos = (truth>0.5).float()
        neg = (truth<0.5).float()
        pos_weight = pos.sum().item() + 1e-12
        neg_weight = neg.sum().item() + 1e-12
        loss = (0.25*pos*loss/pos_weight + 0.75*neg*loss/neg_weight).sum()

    return loss

均衡交叉熵损失

均衡交叉熵损失（balanced corss entropy, BCE）和WCE类似，唯一的不同之处在于也对负样本进行了加权。

如下式所示：

$$\text{BCE}\left(p, \hat{p}\right) = -\left(\beta p \log\left(\hat{p}\right) + (1 - \beta)(1-p) \log\left(1 - \hat{p}\right)\right)$$

Focal loss

在机器学习任务中，除了会遇到严重的类别样本数不均衡问题之外，经常也会遇到容易识别的样本数目和难识别的样本数目不均衡的问题。为了解决这一问题，何凯明大神提出了Focal loss。

Focal loss尝试降低easy example对损失的贡献，这样网络会集中注意力在难样本上。

FL定义如下：

$$\text{FL}\left(p, \hat{p}\right) = -\left(\alpha (1 - \hat{p})^{\gamma} p \log\left(\hat{p}\right) + (1 - \alpha) \hat{p}^{\gamma} (1-p) \log\left(1 - \hat{p}\right)\right)$$

上述公式为二分类问题的Focal loss，可以看出对于每一个样本，使用$(1-\hat p)^\gamma$作为其识别难易程度的指标，预测值$\hat p$越大代表对其进行预测越容易，因而其在总体损失中的占比应该越小。

对于多分类问题，其形式为：

$$FL(p_t)=-\alpha_t(1-p_t)^\gamma \log(p_t)$$

对于每一个样本，$p_t$为模型预测出其属于其真实类别的概率，$\alpha_t$可用于调节不同类别之间的权重。将$\lambda$设置为0便可以得到BCE。

下述代码既可以计算二分类问题，也可以计算多分类问题的focal loss。

class RobustFocalLoss2d(nn.Module):
    #assume top 10% is outliers
    def __init__(self, gamma=2, size_average=True):
        super(RobustFocalLoss2d, self).__init__()
        self.gamma = gamma
        self.size_average = size_average

    def forward(self, logit, target, class_weight=None, type='sigmoid'):
        target = target.view(-1, 1).long()

        if type=='sigmoid':
            if class_weight is None:
                class_weight = [1]*2 # [0.5, 0.5]

            prob = torch.sigmoid(logit)
            prob = prob.view(-1, 1)
            prob = torch.cat((1-prob, prob), 1)
            # select表示样本属于的类别(one-hot)形式
            select = torch.FloatTensor(len(prob), 2).zero_().cuda()
            select.scatter_(1, target, 1.)

        elif type == 'softmax':
            B, C, H, W = logit.size()
            if class_weight is None:
                class_weight =[1]*C #[1/C]*C

            logit = logit.permute(0, 2, 3, 1).contiguous().view(-1, C)
            prob = F.softmax(logit,1)
            select = torch.FloatTensor(len(prob), C).zero_().cuda()
            select.scatter_(1, target, 1.)

        # 各个类别的损失对应的权重
        class_weight = torch.FloatTensor(class_weight).cuda().view(-1,1)
        class_weight = torch.gather(class_weight, 0, target)

        # 样本属于真实类别的概率
        prob = (prob * select).sum(1).view(-1, 1)
        prob = torch.clamp(prob, 1e-8, 1-1e-8)

        focus = torch.pow((1-prob), self.gamma)
        #focus = torch.where(focus < 2.0, focus, torch.zeros(prob.size()).cuda())
        focus = torch.clamp(focus, 0, 2)

        batch_loss = - class_weight * focus * prob.log()

        if self.size_average:
            loss = batch_loss.mean()
        else:
            loss = batch_loss

        return loss

与最近的cell的距离

有文章在交叉熵损失的基础上加入了一个距离函数，是的网络学习两个结束目标的分离边界。如下所示：

$$\text{BCE}\left(p, \hat{p}\right) + w_0\cdot\exp\left(-\frac{(d_1(x) + d_2(x))^2}{2\sigma^2}\right)$$

其中$d_1(x)$和$d_2(x)$为两个距离函数计算点x与最近和第二近的cell的距离。

计算损失函数中的指数项会降低训练的速度，因此一般将距离和输入图片一起输入神经网络。

基于重合度的度量

Dice损失或F1分数

Dice系数

根据Lee Raymond Dice命名，是一种集合相似度度量函数，通常用于计算两个样本的相似度（取值范围为[0,1]）。

Dice系数和Jaccard index类似：

$$\text{DC} = \frac{2 TP}{2 TP + FP + FN} = \frac{2|X \cap Y|}{|X| + |Y|}$$ $$\text{IoU} = \frac{TP}{TP + FP + FN} = \frac{|X \cap Y|}{|X| + |Y| - |X \cap Y|}$$

$|X \bigcap Y|$表示两个样本之间的交集；$|X|$和$|Y|$分别表示X和Y的元素个数，分母的系数为2，用于处理两个集合存在交集的情况。可以看出DC大于IoU。

Dice损失

可以将dice系数定义为损失函数：

$$\text{DL}\left(\mathbf{p}, \mathbf{\hat{p}}\right) = 1 - \frac{2 \langle\mathbf{p}, \mathbf{\hat{p}}\rangle}{\left\lVert\mathbf{p}\right\rVert_2^2 + \left\lVert\mathbf{\hat{p}}\right\rVert_2^2}$$

其中$p\in \{0,1\}^n$，$\hat p \in [0,1]$。实际计算方式如下：

1. 分子的计算

   将$|X \bigcap Y|$近似为预测图和真实标注之间的点乘，点乘的结果进行元素相加：

   $$|X \bigcap Y| = \begin{bmatrix} 0.01 & 0.03 & 0.02 & 0.02 \\ 0.05 & 0.12 & 0.09 & 0.07 \\ 0.89 & 0.85 & 0.88 & 0.91 \\ 0.99 & 0.97 & 0.95 & 0.97 \end{bmatrix} * \begin{bmatrix} 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 \\ 1 & 1 & 1 & 1 \\ 1 & 1 & 1 & 1 \end{bmatrix}$$

接着进行元素相加：

$$|X \bigcap Y| = \begin{bmatrix} 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 \\ 0.89 & 0.85 & 0.88 & 0.91 \\ 0.99 & 0.97 & 0.95 & 0.97 \end{bmatrix} \rightarrow 7.41 (逐元素相加和)$$

1. 分母的计算

   可以直接取元素的平方相加，也可以直接进行元素相加。

   $$|X| = \begin{bmatrix} 0.01 & 0.03 & 0.02 & 0.02 \\ 0.05 & 0.12 & 0.09 & 0.07 \\ 0.89 & 0.85 & 0.88 & 0.91 \\ 0.99 & 0.97 & 0.95 & 0.97 \end{bmatrix} ^{2(可选)} \rightarrow 7.82 (逐元素相加和)$$ $$|Y| = \begin{bmatrix} 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 \\ 1 & 1 & 1 & 1 \\ 1 & 1 & 1 & 1 \end{bmatrix} ^{2(可选)} \rightarrow 8 (逐元素相加和)$$

   Dice损失比较适合样本极度不平衡的情况，但会对反向传播造成影响，导致反向传播不稳定。

Dice损失和交叉熵损失的对比及选择

选择交叉熵损失的原因：

交叉熵损失函数相比于基于重合度度量的损失函数，具有更简单的梯度形式，dice损失的梯度形式比较复杂。

选择Dice损失的原因：

该类损失的真实目标就是最大化预测值和真实值的重合度，更为直接，且一般来说，该损失更适用于类别不均衡问题。

weighted soft dice的代码实现：

def criterion_pixel(logit_pixel, truth_pixel):
    batch_size = len(logit_pixel)
    logit = logit_pixel.view(batch_size,-1)
    truth = truth_pixel.view(batch_size,-1)
    assert(logit.shape==truth.shape)

    loss = soft_dice_criterion(logit, truth)

    loss = loss.mean()
    return loss

def soft_dice_criterion(logit, truth, weight=[0.2,0.8]):

    batch_size = len(logit)
    probability = torch.sigmoid(logit)

    p = probability.view(batch_size,-1)
    t = truth.view(batch_size,-1)
    # 计算正样本和负样本的权重
    w = truth.detach()
    w = w*(weight[1]-weight[0])+weight[0]

    p = w*(p*2-1)  #convert to [0,1] --> [-1, 1]
    t = w*(t*2-1)

    intersection = (p * t).sum(-1)
    union =  (p * p).sum(-1) + (t * t).sum(-1)
    dice  = 1 - 2*intersection/union

    loss = dice
    return loss

Tversky损失

TL损失是DL的改进，该FP和FN加上权重。

$$\text{DL}\left(\mathbf{p}, \mathbf{\hat{p}}\right) = \frac{ \langle\mathbf{p}, \mathbf{\hat{p}}\rangle}{\langle\mathbf{p}, \mathbf{\hat{p}}\rangle + \beta\langle\mathbf{1 - p}, \mathbf{\hat{p}}\rangle + (1 - \beta)\langle\mathbf{p}, \mathbf{1 - \hat{p}}\rangle}$$

将$\beta$设置为$\frac{1}{2}$，则变为DL损失。

交叉熵和Dice损失的结合

也有将交叉熵和Dice损失进行结合的损失函数：

$$Loss = BCE + DiceLoss$$

代码如下：

# reference: https://github.com/asanakoy/kaggle_carvana_segmentation
def dice_loss(preds, trues, weight=None, is_average=True):
    num = preds.size(0)
    preds = preds.view(num, -1)
    trues = trues.view(num, -1)
    if weight is not None:
        w = torch.autograd.Variable(weight).view(num, -1)
        preds = preds * w
        trues = trues * w
    intersection = (preds * trues).sum(1)
    scores = 2. * (intersection + 1) / (preds.sum(1) + trues.sum(1) + 1)

    if is_average:
        score = scores.sum() / num
        return torch.clamp(score, 0., 1.)
    else:
        return scores

def dice_clamp(preds, trues, is_average=True):
    preds = torch.round(preds)
    return dice_loss(preds, trues, is_average=is_average)

class DiceLoss(nn.Module):
    """
    """
    def __init__(self, size_average=True):
        super().__init__()
        self.size_average = size_average

    def forward(self, input, target, weight=None):
        return 1 - dice_loss(F.sigmoid(input), target, weight=weight, is_average=self.size_average)

class BCEDiceLoss(nn.Module):
    def __init__(self, size_average=True):
        super().__init__()
        self.size_average = size_average
        self.dice = DiceLoss(size_average=size_average)

    def forward(self, input, target, weight=None):
        return nn.modules.loss.BCEWithLogitsLoss(size_average=self.size_average, weight=weight)(input, target) + self.dice(input, target, weight=weight)

为了解决样本类别不平衡问题，还可以使用加权BCE+DICE。

class SoftDiceLoss(nn.Module):
    """二分类加权dice损失
    """
    def __init__(self, size_average=True, weight=[0.2, 0.8]):
        """
        weight: 各类别权重
        """
        super(SoftDiceLoss, self).__init__()
        self.size_average = size_average
        self.weight = torch.FloatTensor(weight)
    
    def forward(self, logit_pixel, truth_pixel):
        batch_size = len(logit_pixel)
        logit = logit_pixel.view(batch_size, -1)
        truth = truth_pixel.view(batch_size, -1)
        assert(logit.shape == truth.shape)

        loss = self.soft_dice_criterion(logit, truth)

        if self.size_average:
            loss = loss.mean()
        return loss

    def soft_dice_criterion(self, logit, truth):
        batch_size = len(logit)
        probability = torch.sigmoid(logit)

        p = probability.view(batch_size, -1)
        t = truth.view(batch_size, -1)
        # 向各样本分配所属类别的权重
        w = truth.detach()
        self.weight = self.weight.type_as(logit)
        w = w * (self.weight[1] - self.weight[0]) + self.weight[0]

        p = w * (p*2 - 1)  #convert to [0,1] --> [-1, 1]
        t = w * (t*2 - 1)

        intersection = (p * t).sum(-1)
        union =  (p * p).sum(-1) + (t * t).sum(-1)
        dice  = 1 - 2 * intersection/union

        loss = dice
        return loss


class SoftBCEDiceLoss(nn.Module):
    """加权BCE+DiceLoss
    """
    def __init__(self, size_average=True, weight=[0.2, 0.8]):
        """
        weight: weight[0]为负类的权重，weight[1]为正类的权重
        """
        super(SoftBCEDiceLoss, self).__init__()
        self.size_average = size_average
        self.weight = weight
        self.bce_loss = nn.BCEWithLogitsLoss(size_average=self.size_average, pos_weight=torch.tensor(self.weight[1]/self.weight[0]))
        self.softdiceloss = SoftDiceLoss(size_average=self.size_average, weight=weight)
    
    def forward(self, input, target):
        soft_bce_loss = self.bce_loss(input, target)
        soft_dice_loss = self.softdiceloss(input, target)
        loss = soft_bce_loss + soft_dice_loss

        return loss