【分割】U-Net++

概述

以往的分割模型比如U-Net和FCN效果都很好，其中一个原因就是使用了skip-connection，但是他们使用skip-connection的时候都是直接将浅层和深层的不同语义级别的特征进行强制融合。
本文基于这样一种假设：语义接近的特征层进行融合会使得优化器更容易优化网络，并且提高网络对于细粒度特征的捕捉。提出了一种新的skip-connection的技术，即使用重叠的稠密卷积代替粗暴的特征融合，同时引入深度监督，获得精度和速度的权衡。
主要特点：

• 网络结合了类DenseNet结构，密集的跳跃连接提高了梯度流动性。
• 将U-Net的空心结构填满，连接了编码器和解码器特征图之间的语义鸿沟。
• 使用了深度监督，可以进行剪枝。

细节

skip-connection

这个结构本来的作用是融合浅层细粒度的信息与深层高语义级别的特征，并且这样的做法在之前的研究中证明是有效的。作者重新设计了这个结构，使得encoder和decoder的特征级别尽可能的接近。
大体的结果还是U-Net的结构，中间的skip-connection变复杂了，添加了稠密卷积块，每个稠密卷积块由三个稠密卷积层构成，每一个卷积层前面都有一个连接层，这个连接层融合了当前层前面所有的稠密卷积层的输出和下一层上采样得到的相同尺寸的输出。
并且观察可以发现对于$x^{0,k},k=1,2,3,4$他们都是具有全分辨率的特征图

上面过程的一般化表示如下：
比如$x^{1,0}$它的输入就是上一层的$x^{0,0}$的输出，$x^{0,3}$它的输入就是$x^{0,0}$、$x^{0,1}$、$x^{0,2}$的concat在和下一层$x^{1,2}$上采样得到的相同尺寸的特征图concat之后结果。

deep supervision

深度监督是什么：就是在深度神经网络的某些中间隐藏层加了一个辅助的分类器作为一种网络分支来对主干网络进行监督的技巧，用来解决深度神经网络训练梯度消失和收敛速度过慢等问题。这种辅助的分支分类器能够起到一种判断隐藏层特征图质量好坏的作用，并且就是损失函数不只要监督最后的结果还要监督中间的分支。
使用深度监督评估$x^{0,k},k=1,2,3,4$这四个具有全分辨率的特征图。

首先对于这四个节点，后面都接c（C是类别个数）个1x1的卷积核和一个sigmoid函数，得到最后的分割结果
然后按照以下的损失函数在训练集上训练：

然后在验证集上找到最优的层数，接着在测试集上就可以选择模式了。一种是精确模式，取四个输出特征图的平均就是最后的分割结果；另一种是快速模式，我们在验证集上得到的四个层数中选择一个最优的，用它的输出最为最后的分割结果，一方面相对于准确模式快了很多，另一方面，可以达到模型减枝的效果，如下图所示：

作者本人的回答（U-Net++的科研历程）

作者回答
以下是U-Net的结构，4层是作者在他当时的数据集上得到的一个较好的结果，也就是说，到底几层，是需要实验测试的。

那么作者肯定是做了一系列实验，比如如下四个。

那么问题就出现了，我们能不能让网络自己学习到底选几层呢？
直接把上面的四个网络接到一起，用一个网络表示，那么上面四个U-Net就是这个网络的子集了，那么在训练的时候，根据数据的不同，自己就得到了几层好一点，不需要人为设定了，并且他们可以共用一个特征提取器，而这个特征提取器又可以使用更先进的骨干网络代替。

但是问题来了，这个网络显然无法训练，不会有任何梯度会经过这个红色区域，因为它和算loss function的地方是在反向传播时是断开的。
那么怎么解决呢？

• 第一个是用deep supervision，强行加梯度是吧，关于这个，我待会儿展开来说。
• 第二个解决方案是把结构改成这样子：而这个结构也被人提出来了，但是它没有保留skip-connection的长链。
  
  那么作者的想法就是使用方法1并且长链+短链，于是U-Net++出现了。
  

然后加上了深度监督，发现了意外之喜，可以剪纸！！

简答实现

import paddle
import paddle.nn as nn


# 两次卷积操作
# 卷积计算公式：
# 输出大小 = (输入大小 − Filter + 2Padding )/Stride+1
class VGGBlock(nn.Layer):
    def __init__(self,in_channels,out_channels):
        super(VGGBlock, self).__init__()
        self.layer=nn.Sequential(
            nn.Conv2D(in_channels, out_channels, 3, 1, 1),
            nn.BatchNorm2D(out_channels),
            nn.LeakyReLU(),

            nn.Conv2D(out_channels, out_channels, 3, 1, 1),
            nn.BatchNorm2D(out_channels),
            nn.LeakyReLU()
        )
    def forward(self,x):
        return self.layer(x)

# 将decoder当前层上采样并且和encoder当前层做concat
# 这里不再使用反卷积进行上采样了 而是使用线性插值法
class Up(nn.Layer):
    def __init__(self):
        super(Up, self).__init__()
        self.layer=nn.Sequential(
            nn.UpsamplingBilinear2D(scale_factor=2)
        )

    def forward(self,x1,x2):
        x1=self.layer(x1)
        # 因为tensor是ncwh的 我们需要在c维度上concat 所以axis是1
        return paddle.concat([x2,x1],axis=1)


class UNetPlusPlus(nn.Layer):
    def __init__(self,num_classes=2,deep_supervision=False):
        super(UNetPlusPlus, self).__init__()
        self.num_classes=num_classes
        self.deep_supervision=deep_supervision
        filters=[64, 128, 256, 512, 1024]
        self.pool= nn.MaxPool2D(2)
        self.up=Up()


        self.conv0_0 =VGGBlock(3,filters[0])
        self.conv1_0 =VGGBlock(filters[0],filters[1])
        self.conv2_0 =VGGBlock(filters[1],filters[2])
        self.conv3_0 =VGGBlock(filters[2],filters[3])
        self.conv4_0 =VGGBlock(filters[3],filters[4])

        self.conv0_1 = VGGBlock(filters[0]+filters[1], filters[0])
        self.conv1_1 = VGGBlock(filters[1]+filters[2], filters[1])
        self.conv2_1 = VGGBlock(filters[2]+filters[3], filters[2])
        self.conv3_1 = VGGBlock(filters[3]+filters[4], filters[3])

        self.conv0_2 = VGGBlock(filters[0]*2 + filters[1], filters[0])
        self.conv1_2 = VGGBlock(filters[1]*2 + filters[2], filters[1])
        self.conv2_2 = VGGBlock(filters[2]*2 + filters[3], filters[2])

        self.conv0_3 = VGGBlock(filters[0] * 3 + filters[1], filters[0])
        self.conv1_3 = VGGBlock(filters[1] * 3 + filters[2], filters[1])

        self.conv0_4 = VGGBlock(filters[0] * 4 + filters[1], filters[0])

        if self.deep_supervision:
            self.final1=nn.Conv2D(filters[0],self.num_classes,3,1,1)
            self.final2=nn.Conv2D(filters[0],self.num_classes,3,1,1)
            self.final3=nn.Conv2D(filters[0],self.num_classes,3,1,1)
            self.final4=nn.Conv2D(filters[0],self.num_classes,3,1,1)
        else:
            self.final = nn.Conv2D(filters[0], self.num_classes, 3, 1, 1)



    def forward(self,x):
        x0_0=self.conv0_0(x)
        x1_0=self.conv1_0(self.pool(x0_0))
        x0_1=self.conv0_1(self.up(x1_0,x0_0))

        x2_0 = self.conv2_0(self.pool(x1_0))
        x1_1=self.conv1_1(self.up(x2_0,x1_0))
        x0_2=self.conv0_2(self.up(x1_1,paddle.concat([x0_0,x0_1],axis=1)))

        x3_0 = self.conv3_0(self.pool(x2_0))
        x2_1 = self.conv2_1(self.up(x3_0, x2_0))
        x1_2 = self.conv1_2(self.up(x2_1, paddle.concat([x1_0, x1_1], axis=1)))
        x0_3 = self.conv0_3(self.up(x1_2, paddle.concat([x0_0, x0_1,x0_2], axis=1)))

        x4_0 = self.conv4_0(self.pool(x3_0))
        x3_1 = self.conv3_1(self.up(x4_0, x3_0))
        x2_2 = self.conv2_2(self.up(x3_1, paddle.concat([x2_0, x2_1], axis=1)))
        x1_3 = self.conv1_3(self.up(x2_2, paddle.concat([x1_0, x1_1, x1_2], axis=1)))
        x0_4 = self.conv0_4(self.up(x1_3, paddle.concat([x0_0, x0_1, x0_2,x0_3], axis=1)))

        if self.deep_supervision:
            output1 = self.final1(x0_1)
            output2 = self.final2(x0_2)
            output3 = self.final3(x0_3)
            output4 = self.final4(x0_4)
            return [output1, output2, output3, output4]

        else:
            output = self.final(x0_4)
            return output

if __name__ == '__main__':
    # x=paddle.randn(shape=[2,3,256,256])
    unet=UNetPlusPlus()
    # print(net(x).shape)
    paddle.summary(unet, (1,3,256,256))