这两个损失函数出自《F³Net: Fusion, Feedback and Focus for Salient Object Detection》一文中,用于处理显著性检测(二分割)中小目标的问题。对于传统的BCE Loss,其存在以下三个问题:
- 只是简单的将每个像素求BCE再平均,忽视了目标对象的结构
- 对于小目标而言,整张图像的loss会被背景类所主导,导致难以对前景进行学习
- 对象的边缘位置像素非常容易分类错误,不应该与其他位置像素一样给予相似的权重
而解决方案自然是对不同位置的像素进行加权。具体来说,权重最高的地方应该是地方的边缘位置,离边缘越远则越低。那么怎么"非简单粗暴"的实现这一加权过程呢?文章weighted BCE Loss表达式如下: L w b c e s = − ∑ i = 1 H ∑ j = 1 W ( 1 + γ α i j ) ∑ l = 0 1 1 ( g i j s = l ) log P r ( p i j s = l ∣ Ψ ) ∑ i = 1 H ∑ j = 1 W γ α i j L_{w b c e}^s=-\frac{\sum_{i=1}^H \sum_{j=1}^W\left(1+\gamma \alpha_{i j}\right) \sum_{l=0}^1 \mathbf{1}\left(g_{i j}^s=l\right) \log \mathbf{P r}\left(p_{i j}^s=l \mid \Psi\right)}{\sum_{i=1}^H \sum_{j=1}^W \gamma \alpha_{i j}} Lwbces=−∑i=1H∑j=1Wγαij∑i=1H∑j=1W(1+γαij)∑l=011(gijs=l)logPr(pijs=l∣Ψ) 这里的 α i j \alpha_{i j} αij指的就是(i, j)位置像素的权重。如果不进行加权的话,则相当于 α i j \alpha_{i j} αij恒为1。
公式的其他部分暂且不管,重点看 α i j \alpha_{i j} αij是如何计算的: α i j s = ∣ ∑ m , n ∈ A i j g m n s ∑ m , n ∈ A i j 1 − g i j s ∣ \alpha_{i j}^s=\left|\frac{\sum_{m, n \in A_{i j}} g_{m n}^s}{\sum_{m, n \in A_{i j}} 1}-g_{i j}^s\right| αijs= ∑m,n∈Aij1∑m,n∈Aijgmns−gijs 其中 g i j s g_{i j}^s gijs表示(i, j)位置的真值(1或0,对应前景或背景), A i j A_{i j} Aij表示(i, j)周围的像素。
我们取几个特殊值来讨论。假设
g
m
n
s
g_{m n}^s
gmns均为0,
g
i
j
s
g_{i j}^s
gijs为1,相当于当前像素为前景,而周围像素均为背景,是一个小目标,应该给予高权重。类似的,如果
g
m
n
s
g_{m n}^s
gmns均为0,
g
i
j
s
g_{i j}^s
gijs也为0,表明当前位置与周围位置均为背景,此时对应着一个低权重。对权重进行可视化如下:
可以看到,边缘附近的权重偏高(红色),而远离边缘的像素会被归零。从某种程度可以理解为这是不需要显式输入边缘信息的boundary-aware方法。
接下来看weighted IoU Loss。需要注意的是,IoU这个概念天然适合处理小目标,因此对IoU Loss进行加权单纯是为了,其表达式如下: L wiou s = 1 − ∑ i = 1 H ∑ j = 1 W ( g t i j s ∗ p i j s ) ∗ ( 1 + γ α i j s ) ∑ i = 1 H ∑ j = 1 W ( g t i j s + p i j s − g t i j s ∗ p i j s ) ∗ ( 1 + γ α i j s ) L_{\text {wiou }}^s=1-\frac{\sum_{i=1}^H \sum_{j=1}^W\left(g t_{i j}^s * p_{i j}^s\right) *\left(1+\gamma \alpha_{i j}^s\right)}{\sum_{i=1}^H \sum_{j=1}^W\left(g t_{i j}^s+p_{i j}^s-g t_{i j}^s * p_{i j}^s\right) *\left(1+\gamma \alpha_{i j}^s\right)} Lwiou s=1−∑i=1H∑j=1W(gtijs+pijs−gtijs∗pijs)∗(1+γαijs)∑i=1H∑j=1W(gtijs∗pijs)∗(1+γαijs) 这里需要特别注意的是,分子分母中的 ( 1 + γ α i j s ) \left(1+\gamma \alpha_{i j}^s\right) (1+γαijs)不能直接约掉,因为其并不是一个常数,而是会随(i,j)变化而改变的值。上面这公式的思想依旧为离边缘近的像素对IOU计算的贡献更大。
代码实现如下:
def structure_loss(pred, mask):
weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)
wbce = F.binary_cross_entropy_with_logits(pred, mask, reduce='none')
wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))
pred = torch.sigmoid(pred)
inter = ((pred * mask)*weit).sum(dim=(2, 3))
union = ((pred + mask)*weit).sum(dim=(2, 3))
wiou = 1 - (inter + 1)/(union - inter+1)
return (wbce + wiou).mean()
需要注意的是,这里输入的pred(即网络的输出)不需要预先经过sigmoid处理,因为在函数中已经进行了sigmoid。使用时应防止出现进行重复两次sigmoid的情况。
该损失函数涉及到两个超参数,一个是 α i j \alpha_{i j} αij的加权值 γ \gamma γ。 γ \gamma γ为0,表示不进行任何加权; γ \gamma γ越大,表明网络越依赖于 α i j \alpha_{i j} αij来对不同位置进行加权,不同像素位置的权重差异会更大。在代码中 γ = 5 \gamma=5 γ=5。
另一个则为 A i j A_{i j} Aij中"周围像素"的具体定义,在代码中使用的是F.avg_pool2d进行平均池化。