【深度学习】语义分割损失函数之Lovasz Loss
LovaszLoss，由2017年的论文《》中被提出。它主要适用于语义分割的任务中。本文将详细介绍LovaszLoss的基本概念、思想原理，并提供PyTorch的实现代码，帮助大家去更好的理解和使用。
LovaszLoss，由2017年的论文《The Lovász-Softmax loss: A tractable surrogate for the optimization of the intersection-over-union measure in neural networks》中被提出。它主要适用于语义分割的任务中。
本文将详细介绍LovaszLoss的基本概念、思想原理，并提供PyTorch的实现代码，帮助大家去更好的理解和使用。
Lovasz的基本概念
语义分割的任务效果常常用IOU（Intersection Over Union）来评价，那么很自然的一个想法就是，能不能直接使用IOU来作为损失函数呢？
先看IOU的多种表达方式：
I O U = ∣ X ⋂ Y ∣ ∣ X ⋃ Y ∣ = T P T P + F P + F N \begin{aligned}IOU&=\frac{|X\bigcap Y|}{|X\bigcup Y|}\\&=\frac{TP}{TP+FP+FN}\end{aligned} IOU=∣X⋃Y∣∣X⋂Y∣=TP+FP+FNTP
对比DiceLoss
如果你了解过DiceLoss，那么你一定知道它的定义：
D i c e = 2 ⋅ ∣ X ⋂ Y ∣ ∣ X ∣ + ∣ Y ∣ = 2 ⋅ T P 2 ⋅ T P + F P + F N D i c e L o s s = 1 − D i c e \begin{aligned}Dice&=\frac{2\cdot |X\bigcap Y|}{|X|+|Y|}\\&=\frac{2\cdot TP}{2\cdot TP+FP+FN}\\ DiceLoss&=1-Dice\end{aligned} DiceDiceLoss=∣X∣+∣Y∣2⋅∣X⋂Y∣=2⋅TP+FP+FN2⋅TP=1−Dice
由于Dice是离散的，如果DiceLoss延续 1 − D i c e 1-Dice 1−Dice，并且使用Dice定义的那种计算方式，是不能作为神经网络的优化目标。
因此，我们对Dice以模型输出的概率值进行替代计算，从而使得Dice和DiceLoss都获得连续性。即：在连续值预测（如概率输出）中，通过点乘求和近似交集 ∣ X ⋂ Y ∣ |X\bigcap Y| ∣X⋂Y∣，分母直接取预测值 ∣ X ∣ |X| ∣X∣和真实值 ∣ Y ∣ |Y| ∣Y∣的总和。这种操作天然可导，无需额外处理。
不太清楚DiceLoss的可以参考博文：【深度学习】语义分割损失函数之Dice Loss。
同理，如果对LovaszLoss同样的方式定义：
L o v a s z L o s s = 1 − I O U LovaszLoss=1-IOU LovaszLoss=1−IOU
由于Iou是离散的，如果LovaszLoss延续 1 − I O U 1-IOU 1−IOU，并且使用IOU定义的那种计算方式，是不能作为神经网络的优化目标。
那么，能否也能仿照Dice、DiceLoss的方式获得连续性么？即：
- 分子 ∣ X ⋂ Y ∣ |X\bigcap Y| ∣X⋂Y∣的连续化，就是Dice的分子，这没有问题
- 分母 ∣ X ⋃ Y ∣ |X\bigcup Y| ∣X⋃Y∣的连续化，可以按照 ∣ X ⋃ Y ∣ = ∣ X ∣ + ∣ Y ∣ − ∣ X ⋂ Y ∣ |X\bigcup Y| = |X|+|Y|-|X\bigcap Y| ∣X⋃Y∣=∣X∣+∣Y∣−∣X⋂Y∣的方式，其中前者 ∣ X ∣ + ∣ Y ∣ |X|+|Y| ∣X∣+∣Y∣就是Dice的分母，后者 ∣ X ⋂ Y ∣ |X\bigcap Y| ∣X⋂Y∣就是Dice的分子，这貌似也没有什么问题。
那这种方式可行么？
理论上是可以的，这种方式在2016年的论文Optimizing Intersection-Over-Union in Deep Neural Networks for Image Segmentation中已经被提出且使用到了。
显然，本文要讲述的LovaszLoss，并不是使用的这种方法。
虽然LovaszLoss的论文中没有与其进行比较，但作者在github中说本文的LovaszLoss光滑延拓得到的loss，要比DiceLoss那样简单的光滑化（连续画处理）效果好。
主要原因可能是：IOU分母的非线性复杂性较高。即：交集 ∣ X ⋂ Y ∣ |X\bigcap Y| ∣X⋂Y∣将会同时出现在分子和分母中，这使得分母变得更加复杂，直接连续化会导致梯度计算变得困难，尤其是在优化过程中，梯度可能会变得不稳定。
那么，本文中所提出的更好效果的连续化是怎么获得的呢？
LovaszLoss的连续性
LovaszLoss的连续性表达，从理论上看是比较复杂的，也是比较难理解的。
为了和论文中的符号表示一致，我们重新表达下IOU、LovaszLoss：
I O U = J c = ∣ { y ∗ = c } ⋂ { y ~ = c } ∣ ∣ { y ∗ = c } ⋃ { y ~ = c } ∣ IOU=J_c=\frac{|\{y^*=c\}\bigcap \{\tilde y=c\}|}{|\{y^*=c\}\bigcup \{\tilde y=c\}|} IOU=Jc=∣{y∗=c}⋃{y~=c}∣∣{y∗=c}⋂{y~=c}∣
其中：针对于类别 c c c来说， y ∗ ∈ { 0 , 1 } y^*\in\{0,1\} y∗∈{0,1}是对于类别 c c c的one hot形式表达的target标签结果， y ~ ∈ [ 0 , 1 ] \tilde y\in[0,1] y~∈[0,1]是模型预测结果。
当然，这是针对于多分类情况来说的，如果是二分类的情况， c c c默认就是正样本。
那么，LovaszLoss的表达形式则为：
L o v a s z L o s s = Δ J c = 1 − J c LovaszLoss=\Delta _{J_c}=1-J_c LovaszLoss=ΔJc=1−Jc
重新表达完LovaszLoss之后，接下去需要讨论其连续化的方式。
Submodular Function
上面的LovaszLoss公式没法直接连续化，我们对它进行一些调整，化简成以下的形式：
Δ J c = ∣ M c ∣ ∣ { y ∗ = c } ⋃ M c ∣ \Delta _{J_c}=\frac{|M_c|}{|\{y^*=c\}\bigcup M_c|} ΔJc=∣{y∗=c}⋃Mc∣∣Mc∣
其中，自变量 M c M_c Mc表示类别c时，模型预测错误的集合，定义为：
M c = { y ∗ = c , y ~ ≠ c } ⋃ { y ∗ ≠ c , y ~ = c } M_c=\{y^*=c,\tilde y\neq c\}\bigcup \{y^*\neq c,\tilde y= c\} Mc={y∗=c,y~=c}⋃{y∗=c,y~=c}
解释一下， M c M_c Mc的前者表示target类别是c的，模型预测negative；后者表示target类别不是c的，模型预测positive。即模型对第c类的预测结果和target类别不匹配的部分（模型预测错误的集合），前者 F N FN FN，后者 F P FP FP。
对于自变量 M c M_c Mc的取值范围，一般记为 { 0 , 1 } p \{0,1\}^{p} {0,1}p，特征向量表示法， p p p表示样本数量，0表示某个样本预测成功，1表示某个样本预测错误。
为什么， Δ J c \Delta _{J_c} ΔJc可以这样用 M c M_c Mc来表示的呢？
对于 M c M_c Mc，刚说到：
∣ M c ∣ = ∣ F N + F P ∣ |M_c|=|FN+FP| ∣Mc∣=∣FN+FP∣
而分母：
{ y ∗ = c } ⋃ M c = ( T P + F N ) ⋃ ( F N + F P ) \{y^*=c\}\bigcup M_c=(TP+FN)\bigcup (FN+FP) {y∗=c}⋃Mc=(TP+FN)⋃(FN+FP)
因此，可以化简为：
∣ { y ∗ = c } ⋃ M c ∣ = ∣ T P + F N + F P ∣ |\{y^*=c\}\bigcup M_c|=|TP+FN+FP| ∣{y∗=c}⋃Mc∣=∣TP+FN+FP∣
用比较熟悉的方式来推导：
Δ J c = 1 − J c = 1 − I O U = 1 − T P T P + F P + F N = F P + F N T P + F P + F N \begin{aligned}\Delta _{J_c}&=1-J_c=1-IOU\\&=1-\frac{TP}{TP+FP+FN}\\&=\frac{FP+FN}{TP+FP+FN}\end{aligned} ΔJc=1−Jc=1−IOU=1−TP+FP+FNTP=TP+FP+FNFP+FN
可以看出，两个的表达形式是完全一样的。
为什么要化简成，用 M c M_c Mc作为自变量表达的这个样子呢？
因为，对任意的一个离散函数找到其光滑的延拓很难。但是如果一个离散函数是submodular的，那么就已经有成熟数学工具可以直接，将其做光滑延拓，而且延拓后的函数总是凸的，这样就大大方便了优化。
作者发现，原始的 Δ J c \Delta _{J_c} ΔJc并不是submodular的，而用 M c M_c Mc做自变量表达的 Δ J c \Delta _{J_c} ΔJc却是submodular的。
什么是submodular？
记具有 n n n个元素的集合为 [ n ] = { 1 , 2 , . . . , n } [n]=\{1,2,...,n\} [n]={1,2,...,n}，集合 [ n ] [n] [n]的所有子集对应的集合为 2 [ n ] 2^{[n]} 2[n]。如果一个集函数 f : 2 [ n ] → R f:2^{[n]}\to \mathbb{R} f:2[n]→R是子模的，则对于集合 [ n ] [n] [n]的所有真子集对 T ⊂ S T\subset S T⊂S、 S ⊂ [ n ] S\subset [n] S⊂[n]、 T ⊂ [ n ] T\subset [n] T⊂[n]和集合中的所有元素 i ∈ [ n ] i\in [n] i∈[n]，存在收益递减性质：
f ( T ⋃ { i } ) − f ( T ) ≥ f ( S ⋃ { i } ) − f ( S ) f(T\bigcup \{i\})-f(T)\ge f(S\bigcup \{i\})-f(S) f(T⋃{i})−f(T)≥f(S⋃{i})−f(S)
等价的，如果一个集函数 f : 2 [ n ] → R f:2^{[n]}\to \mathbb{R} f:2[n]→R是子模的，则对于所有集合 S ⊂ [ n ] S\subset [n] S⊂[n]、 T ⊂ [ n ] T\subset [n] T⊂[n]，满足：
f ( S ) + f ( T ) ≥ f ( S ⋃ T ) + f ( S ⋂ T ) f(S)+f(T)\ge f(S\bigcup T)+f(S\bigcap T) f(S)+f(T)≥f(S⋃T)+f(S⋂T)
一般地，集合 S ⊂ [ n ] S\subset [n] S⊂[n]可以通过它的特征向量 X S ∈ H n = { 0 , 1 } n X_S\in H_n=\{0,1\}^n XS∈Hn={0,1}n表示，其中 X S ( i ) = 1 X_S(i)=1 XS(i)=1表示 i ∈ S i\in S i∈S，否则 X S ( i ) = 0 X_S(i)=0 XS(i)=0。则集函数也可以定义在特征向量上 f : H n → R f:H_n\to \mathbb {R} f:Hn→R。
子模性表示边际收益递减的性质：当集合变大时，新增元素对函数值的贡献会减少。
对比下，LovaszLoss公式：
-
f : 2 [ n ] → R f:2^{[n]}\to \mathbb{R} f:2[n]→R，对应 Δ J c \Delta_{J_c} ΔJc
-
子集 S S S，对应 M c M_c Mc，使用特征向量 { 0 , 1 } n \{0,1\}^n {0,1}n表示
Lovasz Extension
上文说到，submodular的函数已经有成熟数学工具可以将其做光滑延拓，该数学工具即为lovasz extension。
lovasz extension的作用包括：
-
将离散的子模函数扩展到连续空间，使其在连续值（如模型的概率输出）上定义，从而可以进行梯度计算和优化
-
如果原始的离散函数是子模的，那么其 Lovasz Extension 是凸的。这种凸性使得优化问题更容易求解，避免了局部最优的问题
给定一个子模函数 f : H n → R f:H_n\to \mathbb R f:Hn→R，对应的Lovasz Extension， f ^ : K n → R \hat f:K_n\to \mathbb R f^:Kn→R定义为：
f ^ ( x ) = ∑ i = 1 n x π i ⋅ g i ( x ) \hat f(x)=\sum_{i=1}^n x_{\pi_i}\cdot g_i(x) f^(x)=i=1∑nxπi⋅gi(x)
其中：
-
f ^ ( x ) \hat f(x) f^(x)：表示光滑延拓后的连续化函数
-
π i \pi_i πi：表示 x 1 , x 2 , . . . x n x_1,x_2,...x_n x1,x2,...xn的从大到小的排序索引，即 1 ≥ x π 1 ≥ x π 2 ≥ ⋯ ≥ x π n ≥ 0 1\geq x_{\pi_1} \geq x_{\pi_2} \geq \dots \geq x_{\pi_n}\geq 0 1≥xπ1≥xπ2≥⋯≥xπn≥0
-
g i ( x ) g_i(x) gi(x)：表示Lovasz延拓的梯度，它的定义：
g i ( x ) = f ( X i ) − f ( X i − 1 ) \begin{aligned}g_i(x)=f(X_i)-f(X_{i-1})\end{aligned} gi(x)=f(Xi)−f(Xi−1)
-
X i X_i Xi：与 e π i e_{\pi_i} eπi相关的计算量， X 0 = 0 n X_0=0_n X0=0n、 X i = X i − 1 + e π i X_i=X_{i-1}+e_{\pi_i} Xi=Xi−1+eπi
-
e π i e_{\pi_i} eπi：与 π i \pi_i πi相关的计算量，位置 π i \pi_i πi为1，其他位置为0的向量
注意：这时候新函数 f ^ \hat f f^的定义域已经从离散的 { 0 , 1 } p \{0,1\}^p {0,1}p变到了连续的 [ 0 , 1 ] p [0,1]^p [0,1]p， 此时已经是连续、piecewise linear的，能直接对 x x x 求导，而且导数很简洁，就是 g ( x ) g(x) g(x)。
Lovasz Extension的定义看着就超级复杂且麻烦，并且很不容易理解。
幸好，下面论文里面对其将离散函数映射到连续空间的过程进行梳理。只需要按照这个步骤进行就可以完成了：
-
排序：计算模型预测和taerget的误差 m m m，对其按降序排列，并得到排序索引 π 1 , π 2 , . . . , π p \pi_1,\pi_2,...,\pi_p π1,π2,...,πp
-
梯度计算：延拓后的梯度仅与排序位置相关，这部分的计算过程是最麻烦的，但是论文里面介绍了详细算法：
g i ( m ) = Δ ( π 1 , … , π i ) − Δ ( π 1 , … , π i − 1 ) g_i(m) = \Delta({{\pi_1}, \dots, {\pi_i}}) - \Delta({{\pi_1}, \dots, {\pi_{i-1}}}) gi(m)=Δ(π1,…,πi)−Δ(π1,…,πi−1)
- 线性插值：基于排序后的误差，构造拥有连续性的损失函数：
Δ ˉ = ∑ i = 1 p m i ⋅ g i ( m ) \bar \Delta = \sum_{i=1}^p m_i\cdot g_i(m) Δˉ=i=1∑pmi⋅gi(m)
代码实战
二分类问题
class BinaryLovaszLoss(nn.Module):
"""
二分类BinaryLovaszLoss
"""
def __init__(self):
super(BinaryLovaszLoss, self).__init__()
def forward(self, pred, target):
"""
pred: 模型的输出, 经过sigmoid, 形状为 [B, H, W] (批次大小、图像高度、图像宽度)
target: 标签, 形状为 [B, H, W], 取值范围为0或1
"""
loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels))
return loss
def flatten_binary_scores(scores, labels):
"""展平"""
scores = scores.view(-1)
labels = labels.view(-1)
return scores, labels
def lovasz_hinge_flat(logits, labels):
"""
pred: 模型的输出, 经过 sigmoid, 展平为 [P, C]
target: 标签, 展平为 [P], 取值范围为0或1
"""
if len(labels) == 0:
# only void pixels, the gradients should be 0
return logits.sum() * 0.
# 注意下此时的预测误差，当label为1时，signs为1，当label为0时，signs为-1
signs = 2. * labels.float() - 1.
# 此时label为1时，error=1-输出概率；当label为0时，error=1+输出概率
# 这边的计算误差的方式比较奇怪
errors = (1. - logits * Variable(signs))
errors_sorted, perm = torch.sort(errors, dim=0, descending=True)
perm = perm.data
gt_sorted = labels[perm]
grad = lovasz_grad(gt_sorted)
loss = torch.dot(F.relu(errors_sorted), Variable(grad))
return loss
def lovasz_grad(gt_sorted):
"""计算梯度"""
p = len(gt_sorted)
gts = gt_sorted.sum()
intersection = gts - gt_sorted.float().cumsum(0)
union = gts + (1 - gt_sorted).float().cumsum(0)
jaccard = 1.0 - intersection / union
if p > 1: # cover 1-pixel case
jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
return jaccard
多分类问题：采用one-hot编码
class LovaszLoss(nn.Module):
"""
多分类LovaszLoss
"""
def __init__(self):
super(LovaszLoss, self).__init__()
def forward(self, pred, target):
"""
pred: 模型的输出, 未经过 Softmax, 形状为 [B, C, H, W] (批次大小、类别数、图像高度、图像宽度)
target: 标签, 形状为 [B, H, W], 取值范围为0到C-1
"""
prob = F.softmax(pred, dim=1)
loss_lovasz = lovasz_softmax_flat(
*flatten_probs(prob, target, self.ignore_index)
)
return loss_lovasz
def lovasz_softmax_flat(probs, labels):
"""
pred: 模型的输出, 经过 Softmax, 展平为 [P, C]
target: 标签, 展平为 [P], 取值范围为0到C-1
"""
if probs.numel() == 0:
# only void pixels, the gradients should be 0
return probs.sum() * 0.0
C = probs.size(1)
losses = []
for c in list(range(C)):
fg = (labels == c).float()
if fg.sum() == 0:
continue
class_pred = probs[:, c]
errors = (fg - class_pred).abs()
errors_sorted, perm = torch.sort(errors, 0, descending=True)
fg_sorted = fg[perm]
losses.append(torch.dot(errors_sorted, lovasz_grad(fg_sorted)))
return mean(losses)
def flatten_probs(probs, labels):
"""展平"""
B, C, H, W = probs.size()
# permute：将通道维度 C 移到最后，形状变为 (B, H, W, C)
# contiguous：确保张量在内存中是连续的，便于后续操作
# view：将张量展平为二维 (P, C)，其中 P = B * H * W 是像素总数，C 是类别数
probs = probs.permute(0, 2, 3, 1).contiguous().view(-1, C)
# 将 labels 展平为一维 (P,)，与展平后的 probs 对应
labels = labels.view(-1)
return probs, labels
def lovasz_grad(gt_sorted):
"""计算梯度"""
p = len(gt_sorted)
gts = gt_sorted.sum()
intersection = gts - gt_sorted.float().cumsum(0)
union = gts + (1 - gt_sorted).float().cumsum(0)
jaccard = 1.0 - intersection / union
if p > 1: # cover 1-pixel case
jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
return jaccard
需要注意的是，这边的代码写法应该与输入的shape、参数的shape有关系的，需要针对于具体的情形进行适配和修改。
相关阅读
魔乐社区（Modelers.cn) 是一个中立、公益的人工智能社区，提供人工智能工具、模型、数据的托管、展示与应用协同服务，为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作，由全产业链共同建设、共同运营、共同享有，推动国产AI生态繁荣发展。
更多推荐
12
0- 0
已为社区贡献6条内容
所有评论(0)