简单理解小目标分割中的weighted BCE Loss与weighted IoU Loss

这两个损失函数出自《F³Net: Fusion, Feedback and Focus for Salient Object Detection》一文中，用于处理显著性检测(二分割)中小目标的问题。对于传统的BCE Loss，其存在以下三个问题：

• 只是简单的将每个像素求BCE再平均，忽视了目标对象的结构
• 对于小目标而言，整张图像的loss会被背景类所主导，导致难以对前景进行学习
• 对象的边缘位置像素非常容易分类错误，不应该与其他位置像素一样给予相似的权重

而解决方案自然是对不同位置的像素进行加权。具体来说，权重最高的地方应该是地方的边缘位置，离边缘越远则越低。那么怎么"非简单粗暴"的实现这一加权过程呢？文章weighted BCE Loss表达式如下：$$L_{wbce}^s=-\frac{{\sum }_{i=1}^H{\sum }_{j=1}^W\left(1+\gamma {\alpha }_{ij}\right){\sum }_{l=0}^{1}1\left(g_{ij}^s=l\right)\log \mathbf{Pr}\left(p_{ij}^s=l\mid \Psi \right)}{{\sum }_{i=1}^H{\sum }_{j=1}^W\gamma {\alpha }_{ij}}$$ 这里的${\alpha }_{ij}$指的就是(i, j)位置像素的权重。如果不进行加权的话，则相当于${\alpha }_{ij}$恒为1。

公式的其他部分暂且不管，重点看${\alpha }_{ij}$是如何计算的：$${\alpha }_{ij}^s=\left|\frac{{\sum }_{m,n\in A_{ij}}{g}_{mn}^s}{{\sum }_{m,n\in A_{ij}}1}-g_{ij}^s\right|$$ 其中$g_{ij}^s$表示(i, j)位置的真值(1或0，对应前景或背景)，$A_{ij}$表示(i, j)周围的像素。

我们取几个特殊值来讨论。假设$g_{mn}^s$均为0，$g_{ij}^s$为1，相当于当前像素为前景，而周围像素均为背景，是一个小目标，应该给予高权重。类似的，如果$g_{mn}^s$均为0，$g_{ij}^s$也为0，表明当前位置与周围位置均为背景，此时对应着一个低权重。对权重进行可视化如下：

可以看到，边缘附近的权重偏高(红色)，而远离边缘的像素会被归零。从某种程度可以理解为这是不需要显式输入边缘信息的boundary-aware方法。

接下来看weighted IoU Loss。需要注意的是，IoU这个概念天然适合处理小目标，因此对IoU Loss进行加权单纯是为了统一概念，其表达式如下：$$L_{\text{wiou }}^s=1-\frac{{\sum }_{i=1}^H{\sum }_{j=1}^W\left(g{t}_{ij}^s\ast p_{ij}^s\right)\ast \left(1+\gamma {\alpha }_{ij}^s\right)}{{\sum }_{i=1}^H{\sum }_{j=1}^W\left(g{t}_{ij}^s+p_{ij}^s-g{t}_{ij}^s\ast p_{ij}^s\right)\ast \left(1+\gamma {\alpha }_{ij}^s\right)}$$ 这里需要特别注意的是，分子分母中的$\left(1+\gamma {\alpha }_{ij}^s\right)$不能直接约掉，因为其并不是一个常数，而是会随(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的情况。

该损失函数涉及到两个超参数，一个是${\alpha }_{ij}$的加权值$\gamma$。$\gamma$为0，表示不进行任何加权；$\gamma$越大，表明网络越依赖于${\alpha }_{ij}$来对不同位置进行加权，不同像素位置的权重差异会更大。在代码中$\gamma =5$。

另一个则为$A_{ij}$中"周围像素"的具体定义，在代码中使用的是F.avg_pool2d进行平均池化。