图像分割的损失函数
基于分布相似度的损失
在此类损失函数中,主要使用信息论中的交叉熵机制来度量模型输出和真实目标之间所包含的信息的相似性。
交叉熵损失
假设$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$
交叉熵损失函数的定义形式如下:
上述交叉熵损失为二维交叉熵损失。
加权交叉熵损失
加权交叉熵损失(weighted corss entropy, WCE)是交叉熵损失的一种变体,其中,所有的正样本都被乘以一个系数以进行加权。该损失函数常用于类别不平衡问题中。例如,当你有一张有10%的黑像素和90%的白像素的图片时,常规的CE效果不会太好。
WCE定义如下式:
减少假负样本的数目,将$\beta$设置为大于1.增加假正样本的数目,将$\beta$设置为小于1。
weighted bce代码实现:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16def 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类似,唯一的不同之处在于也对负样本进行了加权。
如下式所示:
Focal loss
在机器学习任务中,除了会遇到严重的类别样本数不均衡问题之外,经常也会遇到容易识别的样本数目和难识别的样本数目不均衡的问题。为了解决这一问题,何凯明大神提出了Focal loss。
Focal loss尝试降低easy example对损失的贡献,这样网络会集中注意力在难样本上。
FL定义如下:
上述公式为二分类问题的Focal loss,可以看出对于每一个样本,使用$(1-\hat p)^\gamma$作为其识别难易程度的指标,预测值$\hat p$越大代表对其进行预测越容易,因而其在总体损失中的占比应该越小。
对于多分类问题,其形式为:
对于每一个样本,$p_t$为模型预测出其属于其真实类别的概率,$\alpha_t$可用于调节不同类别之间的权重。将$\lambda$设置为0便可以得到BCE。
下述代码既可以计算二分类问题,也可以计算多分类问题的focal loss。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51class 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的距离
有文章在交叉熵损失的基础上加入了一个距离函数,是的网络学习两个结束目标的分离边界。如下所示:
其中$d_1(x)$和$d_2(x)$为两个距离函数计算点x与最近和第二近的cell的距离。
计算损失函数中的指数项会降低训练的速度,因此一般将距离和输入图片一起输入神经网络。
基于重合度的度量
Dice损失或F1分数
Dice系数
根据Lee Raymond Dice命名,是一种集合相似度度量函数,通常用于计算两个样本的相似度(取值范围为[0,1])。
Dice系数和Jaccard index类似:
$|X \bigcap Y|$表示两个样本之间的交集;$|X|$和$|Y|$分别表示X和Y的元素个数,分母的系数为2,用于处理两个集合存在交集的情况。可以看出DC大于IoU。
Dice损失
可以将dice系数定义为损失函数:
其中$p\in \{0,1\}^n$,$\hat p \in [0,1]$。实际计算方式如下:
分子的计算
将$|X \bigcap Y|$近似为预测图和真实标注之间的点乘,点乘的结果进行元素相加:
接着进行元素相加:
分母的计算
可以直接取元素的平方相加,也可以直接进行元素相加。
Dice损失比较适合样本极度不平衡的情况,但会对反向传播造成影响,导致反向传播不稳定。
Dice损失和交叉熵损失的对比及选择
选择交叉熵损失的原因:
交叉熵损失函数相比于基于重合度度量的损失函数,具有更简单的梯度形式,dice损失的梯度形式比较复杂。
选择Dice损失的原因:
该类损失的真实目标就是最大化预测值和真实值的重合度,更为直接,且一般来说,该损失更适用于类别不均衡问题。
weighted soft dice的代码实现:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31def 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加上权重。
将$\beta$设置为$\frac{1}{2}$,则变为DL损失。
交叉熵和Dice损失的结合
也有将交叉熵和Dice损失进行结合的损失函数:
代码如下:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40# 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。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64class 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