U^2 Net显著性检测分割抠图

论文：U^2 -Net: Going Deeper with Nested U-Structure for Salient Object Detection

代码：https://github.com/xuebinqin/U-2-Net

1. 简介

这是一篇关于显著性检测的文章，作者是秦雪彬大佬，CVPR2019 BASNet的作者。文章是关于显著性检测的，但其效果突出被用于分割、抠图等领域，也取得了卓越的效果，详情参见GitHub 不断有新的应用出现。文章提出了一种两级嵌套的U-Net结构的网络：U^2 Net。该网络优点：

• 1.提出新的RSU(ReSidual U-blocks)模块，融合不同尺度感受野的特征，来捕捉不同尺度的上下文信息；
• 2.基于RSU模块的池化(pooling) 操作，在不显著增加计算成本的前提下，增加了整个网络结构的深度(depth).

因为这些设计使得我们可以从头训练网络，而不需要使用现有的图像分类网络backbone。

ps: 因为这篇论文模型结构比较简洁明了，开源的代码也比较清晰，所以针对U^2 Net的介绍会更侧重于使用代码讲解。

2. 方法设计

2.1 RSU-L

上图为普通卷积block，Res-like block，Inception-like block，Dense-like block和Residual U-blocks的对比图，明显可以看出Residual U-blocks类似一个简单的U-Net。**其中L是编码器中的层数，Cin，Cout表示输入和输出通道，M表示RSU内部层中的通道数。**采用的是encoder-decoder的结构，下采样使用了大量的池化层，解码上采样使用的是双线性插值。

文章定义了卷积模块：

class REBNCONV(nn.Module):
    def __init__(self,in_ch=3,out_ch=3,dirate=1):
        super(REBNCONV,self).__init__()

        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
        self.bn_s1 = nn.BatchNorm2d(out_ch)
        self.relu_s1 = nn.ReLU(inplace=True)

    def forward(self,x):

        hx = x
        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))

        return xout

主要注意膨胀系数：dirate，dirate==1的时候就是普通卷积，dirate!=1的时候是空洞卷积。

上图e中显示的就是RSU-7, 下面以RSU-7为例进行说明，其他层的代码几乎一致：

### RSU-7 ###
class RSU7(nn.Module):#UNet07DRES(nn.Module):

    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
        super(RSU7,self).__init__()

        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)

        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)

        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2) ###

        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)

    def forward(self,x):

        hx = x
        hxin = self.rebnconvin(hx)

        hx1 = self.rebnconv1(hxin)
        hx = self.pool1(hx1)

        hx2 = self.rebnconv2(hx)
        hx = self.pool2(hx2)

        hx3 = self.rebnconv3(hx)
        hx = self.pool3(hx3)

        hx4 = self.rebnconv4(hx)
        hx = self.pool4(hx4)

        hx5 = self.rebnconv5(hx)
        hx = self.pool5(hx5)

        hx6 = self.rebnconv6(hx)

        hx7 = self.rebnconv7(hx6)

        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1)) # 残差连接 F1(x)+U(F1(x))
        hx6dup = _upsample_like(hx6d,hx5) # 上采样 双线性插值

        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
        hx5dup = _upsample_like(hx5d,hx4)

        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
        hx4dup = _upsample_like(hx4d,hx3)

        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
        hx3dup = _upsample_like(hx3d,hx2)

        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
        hx2dup = _upsample_like(hx2d,hx1)

        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))

        return hx1d + hxin

对应结构的本质如下：

十分类似残差块，将第二个weight layer换成了U-net。这种设计使得网络能从多个尺度直接从残差块中提取特征。并且U结构的计算开销很小，因为大多数操作都是在下采样的特征映射上进行的。详细展开来讲RSU主要由三部分组成：

• 一个输入卷积层，它将输入的feature map x $(H\times W\times C_{in})$转换成中间feature map $F1(x)$，$F1(x)$通道数为$C_{out}$。这是一个用于局部特征提取的普通卷积层。
• 一个U-like的对称的encoder-decoder结构，高度为L，以中间feature map $F1(x)$为输入，去学习提取和编码多尺度文本信息$U(F1(x))$,U表示类U-Net结构。更大L会得到更深层的U-block(RSU)，更多的池操作，更大的感受野和更丰富的局部和全局特征。配置此参数允许从具有任意空间分辨率的输入特征图中提取多尺度特征。从逐渐降采样特征映射中提取多尺度特征，并通过渐进上采样、合并和卷积等方法将其编码到高分辨率的特征图中。这一过程减少了大尺度直接上采样造成的细节损失。
• 一种残差连接，它通过求和来融合局部特征和多尺度特征：$F1(x)+U(F1(x))$

2.2 U^2 net网络结构

U^2 Net整体结构如上图所示，整个结构较为清晰，En_1和De_1，En_2和De_2，En_3和De_3，En_,4和De_4，En_,5和De_5以及En_6分别使用了RSU-7，RSU-6，RSU-5，RSU-4，RSU-4F和RSU-4F。图中红字基本标识出来了。

文章给后续设计留下了空间，1. U^2-Net的每一个Block都是一个U-Net结构的模块，即上述Residual U-blocks。当然，你也可以继续Going Deeper, 每个Block里面的U-Net的子Block仍然可以是一个U-Net结构，命名为U^3-Net。2. 更改不同的RSU-L的L, 3. 更改每个RSU中的I,M,O 通道数，文章据此还提出了轻量化网络U^2 Netp，模型体积很小，但是拥有不俗的效果。

##### U^2-Net ####
class U2NET(nn.Module):

    def __init__(self,in_ch=3,out_ch=1):
        super(U2NET,self).__init__()

        self.stage1 = RSU7(in_ch,32,64)
        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.stage2 = RSU6(64,32,128)
        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.stage3 = RSU5(128,64,256)
        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.stage4 = RSU4(256,128,512)
        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.stage5 = RSU4F(512,256,512)
        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)

        self.stage6 = RSU4F(512,256,512)

        # decoder
        self.stage5d = RSU4F(1024,256,512)
        self.stage4d = RSU4(1024,128,256)
        self.stage3d = RSU5(512,64,128)
        self.stage2d = RSU6(256,32,64)
        self.stage1d = RSU7(128,16,64)

        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)

        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)

    def forward(self,x):

        hx = x

        #stage 1
        hx1 = self.stage1(hx)
        hx = self.pool12(hx1)

        #stage 2
        hx2 = self.stage2(hx)
        hx = self.pool23(hx2)

        #stage 3
        hx3 = self.stage3(hx)
        hx = self.pool34(hx3)

        #stage 4
        hx4 = self.stage4(hx)
        hx = self.pool45(hx4)

        #stage 5
        hx5 = self.stage5(hx)
        hx = self.pool56(hx5)

        #stage 6
        hx6 = self.stage6(hx)
        hx6up = _upsample_like(hx6,hx5)

        #-------------------- decoder --------------------
        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
        # 类似FPN。每个block的输出结果和上一个（下一个block）结果做融合（cat），然后输出。
        hx5dup = _upsample_like(hx5d,hx4)
        # 由于每个block做了下采样，为了resize到原图，需要做一个上采样， 文章用的双线性插值

        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
        hx4dup = _upsample_like(hx4d,hx3)

        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
        hx3dup = _upsample_like(hx3d,hx2)

        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
        hx2dup = _upsample_like(hx2d,hx1)

        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))


        #side output
        d1 = self.side1(hx1d)
        # 把每一个block输出结果，转换成WxHx1的mask最后过一个sigmod就可以得到每个block输出的概率图。

        d2 = self.side2(hx2d)
        d2 = _upsample_like(d2,d1)

        d3 = self.side3(hx3d)
        d3 = _upsample_like(d3,d1)

        d4 = self.side4(hx4d)
        d4 = _upsample_like(d4,d1)

        d5 = self.side5(hx5d)
        d5 = _upsample_like(d5,d1)

        d6 = self.side6(hx6)
        d6 = _upsample_like(d6,d1)

        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
        # 6个block cat一起之后做特征融合，然后再做输出，结果就是d0(单通道)的结果，其他的输出都是为了计算loss

        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)

2.3 损失函数设计

bce_loss = nn.BCELoss(size_average=True) # Binary Cross Entropy

def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v):

	loss0 = bce_loss(d0,labels_v)
	loss1 = bce_loss(d1,labels_v)
	loss2 = bce_loss(d2,labels_v)
	loss3 = bce_loss(d3,labels_v)
	loss4 = bce_loss(d4,labels_v)
	loss5 = bce_loss(d5,labels_v)
	loss6 = bce_loss(d6,labels_v)

	loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
	print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n"%(loss0.data.item(),loss1.data.item(),loss2.data.item(),loss3.data.item(),loss4.data.item(),loss5.data.item(),loss6.data.item()))

	return loss0, loss

类似于HED算法的deep supervision方式，作者设计了如下函数：

其中，M=6, 为U2Net 的 Sup1, Sup2, …, Sup6 stage. $w_{side}^{(m)}{l}^{(m)}side$为对应的损失函数输出和权重；$w_{fuse}{l}_{fuse}$为融合的损失函数和权重;对于每一个$l$使用的都是标准的BCE Loss：

从代码上看很简单，分别计算$l_0,l_1,\dots l_6$的BCE Loss, 然后全部相加，最后返回$loss0,loss$

这里使用BCE Loss的原因是使用交叉熵作为损失函数后，反向传播的梯度不在于sigmoid函数的导数有关了。这就从一定程度上避免了梯度消失。[2]

Reference:

[1] Qin X, Zhang Z, Huang C, et al. U2-Net: Going deeper with nested U-structure for salient object detection[J]. Pattern Recognition, 2020, 106: 107404.

[2] https://blog.csdn.net/geter_CS/article/details/84747670