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

论文

1 Abstract

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

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

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

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

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

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

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

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

; 2 Introduction

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

​ 文章的五个贡献:

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

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

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

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

​ ⑤ 展现了可扩展性。

3 Backbone

3.1 Motivation

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

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

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

; 3.2 Structure

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

3.3 Deep supervision

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

; 3.4 Model pruning

UNet++学习笔记(主干网络+代码)
  1. 集成模式,其中收集所有分割分支的分割结果,然后取其平均值;
  2. 剪枝模式,分割分支,其选择决定了模型修剪的程度和速度增益,例如上图。

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

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

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

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

; 代码

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

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

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])

        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

Original: https://blog.csdn.net/Leon_____/article/details/123311170
Author: 别码了W哥
Title: UNet++学习笔记(主干网络+代码)

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/519089/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球