UNet++学习笔记(主干网络+代码)

论文

1 Abstract

​ 文章提出,UNet主要有以下两大缺陷:

​ ① 网络最优的深度未知,需要通过大量的实验以及集成不同深度的网络,效率低;

​ ② skip connection引入了不必要的限制,即限制仅在相同的尺度进行特征融合。

​ 对此,UNet++进行了以下的优化:

​ ① 利用不同深度UNet的有效集成(这些UNet共享一个编码器),通过监督学习来搜索最优深度;

​ ② 重新设计skip connection,使得解码器的子网络可以聚合不同尺度的特征,更加灵活;

​ ③ 利用剪纸技术来提高UNet++的推理速度。

UNet++学习笔记(主干网络+代码)_第1张图片

2 Introduction

​ 传统的编码器解码器结构 + skip connection结构可以很好的应用于语义分割任务,原因是:其将编码器子网中的浅层细粒度信息与解码器子网中的深层粗粒度信息进行相结合。

​ 文章的五个贡献:

​ ① UNet++内嵌了不同深度的UNet,从而不再是固定的深度结构;

​ ② 更加灵活的skip connection结构,不再是仅融合同一尺度的特征;

​ ③ 设计了一个剪枝操作加快推理速度;

​ ④ 同时训练内嵌的不同深度的UNet引发了UNet之间的协同训练,带来了更好的性能;

​ ⑤ 展现了可扩展性。

3 Backbone

3.1 Motivation

​ 实验发现,更深的UNet不一定更好,因此进行了多组的消融实验。

UNet++学习笔记(主干网络+代码)_第2张图片

​ 在UNete中,需要同时对X01,X02,X03和X04赋予损失函数,从而让内嵌的UNet可以回传梯度。在UNet+到UNet++的过程中,从短连接到长连接,更加有效地利用了多种特征。

3.2 Structure

UNet++学习笔记(主干网络+代码)_第3张图片

3.3 Deep supervision

UNet++学习笔记(主干网络+代码)_第4张图片

3.4 Model pruning

UNet++学习笔记(主干网络+代码)_第5张图片

  1. 集成模式,其中收集所有分割分支的分割结果,然后取其平均值;
  2. 剪枝模式,分割分支,其选择决定了模型修剪的程度和速度增益,例如上图。

​ 以下参考:研习U-Net - 知乎 (zhihu.com)

UNet++学习笔记(主干网络+代码)_第6张图片

UNet++学习笔记(主干网络+代码)_第7张图片

image-20220305211225422

代码

UNet++学习笔记(主干网络+代码)_第8张图片

# 基本的块网络,用于堆叠形成每一个卷积块
class VGGBlock(nn.Module):
    def __init__(self, in_channels, middle_channels, out_channels):
        super().__init__()
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, middle_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(middle_channels)
        self.conv2 = nn.Conv2d(middle_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        return out
    
# UNet++骨干网络
class NestedUNet(nn.Module):
    def __init__(self, num_classes, input_channels=3, deep_supervision=False, **kwargs):
        super().__init__()

        nb_filter = [32, 64, 128, 256, 512]

        self.deep_supervision = deep_supervision
        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
		
        # 第一斜列(左上到右下)
        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])
		
       	# 第二斜列
        self.conv0_1 = VGGBlock(nb_filter[0] * 1 + nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_1 = VGGBlock(nb_filter[1] * 1 + nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv2_1 = VGGBlock(nb_filter[2] * 1 + nb_filter[3], nb_filter[2], nb_filter[2])
        self.conv3_1 = VGGBlock(nb_filter[3] * 1 + nb_filter[4], nb_filter[3], nb_filter[3])
		
        # 第三斜列
        self.conv0_2 = VGGBlock(nb_filter[0] * 2 + nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_2 = VGGBlock(nb_filter[1] * 2 + nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv2_2 = VGGBlock(nb_filter[2] * 2 + nb_filter[3], nb_filter[2], nb_filter[2])
		
        # 第四斜列
        self.conv0_3 = VGGBlock(nb_filter[0] * 3 + nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_3 = VGGBlock(nb_filter[1] * 3 + nb_filter[2], nb_filter[1], nb_filter[1])

        # 第五斜列
        self.conv0_4 = VGGBlock(nb_filter[0] * 4 + nb_filter[1], nb_filter[0], nb_filter[0])
		
        # 1×1卷积核
        if self.deep_supervision:
            self.final1 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final2 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final3 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
        else:
            self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=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(torch.cat([x0_0, self.up(x1_0)], 1))

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

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

        x4_0 = self.conv4_0(self.pool(x3_0))
        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 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

你可能感兴趣的:(网络,深度学习,计算机视觉)