通俗易懂理解U-Net语义分割模型

花花少年

已于 2024-01-24 23:45:31 修改

阅读量7.9k

点赞数 47

分类专栏：深度学习文章标签： U-Net 语义分割

于 2024-01-22 11:55:43 首次发布

本文链接：https://blog.csdn.net/m0_37605642/article/details/135744564

版权

深度学习专栏收录该内容

135 篇文章

订阅专栏

本文详细介绍了U-Net网络，一种经典的语义分割模型，最初用于医疗影像分割，后来广泛应用于其他任务。文章覆盖了U-Net的结构、工作流程、跳跃连接的重要性以及在PyTorch中的实现，包括通用版的改进和相关参考文献。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

温故而知新，可以为师矣！

一、参考资料

语义分割网络U-net介绍

深度学习笔记（二十三）Semantic Segmentation(FCN/U-Net/PSPNet/SegNet/U-Net++/ICNet/DFANet/Fast-SCNN)

图像分割算法U-net

Unet学习

二、U-Net网络相关介绍

U-Net是经典的语义分割网络，它与2015年提出，最初应用在医疗影像分割任务上，由于效果很好，之后被广泛应用在各种分割任务中。

U-Net结构稳定，是典型的 Encoder-Decoder 结构，Encoder进行特征提取，Decoder进行上采样。在数据集较小的时候，推荐使用。

1. U-Net网络结构

U-Net网络由编码器和解码器组成，中间有一个跳跃连接，形状呈U型，所以称为U-Net。U-Net网络结构如下图所示。

Encoder编码器：用于下采样和特征提取，由两个 3x3 的卷积层（no padding）+ReLU+ 2x2 的 max pooling 层（stride=2）反复组成。每次下采样后，输出特征图尺寸减半，通道数翻倍；
Decoder解码器，用于图像尺寸还原及分割，由一个 2x2 的转置卷积层+ReLU+2个 3x3 的卷积层+ReLU+反复构成。
跳跃连接(Skip-connect)：裁剪(crop)编码器对应层的特征图，然后与解码器对应层的特征图进行拼接(concat)。

在这里插入图片描述

2. U-Net网络流程

具体流程如下：

第一层处理
- 输入一张 572×572×1 的图片；
- 使用 64 个 3×3 的卷积核进行卷积，并通过 ReLU 函数得到 64 个 570×570×1 的特征通道；
- 再使用 64 个 3×3 的卷积核进行卷积，并通过 ReLU 函数得到 64 个 568×568×1 的特征通道，即第一层的处理结果。
下采样过程
- 对第一层的处理结果进行 2×2 的最大池化操作，将图片下采样为原来大小的一半：284×284×64；
- 使用 128 个卷积核进一步提取特征，得到一个新的特征图；
- 重复以上步骤，对新的特征图进行下采样，每一层都会经过两次卷积来提取图像特征；
- 每下采样一层，输出特征图尺寸减小一半，卷积核数目增加一倍；
- 最终下采样部分的结果是 28×28×1024，即一共有 1024 个特征层，每一层的特征大小为 28×28。
上采样过程
- 从最右下角开始，把 28×28×1024 的特征矩阵经过512个 2×2 的卷积核进行转置卷积，把矩阵扩大为 56×56×512；
- 为了减少数据丢失，采用把左边降采样时的图片裁剪成相同大小后直接拼接的方法增加特征层（这里是左半边白色部分的512个特征通道），再进行卷积来提取特征；
- 每一层都会进行两次卷积来提取特征，每上采样一层，输出特征图尺寸扩大一倍，卷积核数目减少一半；
- 右边部分从下往上则是4次上采样过程；
- 在最后一步中，选择了2个 1×1 的卷积核把64个特征通道变成2个，也就是最后的 388×388×2，这里是一个二分类的操作，把图片分成背景和前景目标两个类别。

3. 跳跃连接(Skip-connect)

在U-Net网络中，跳跃连接(Skip-connect)实现了特征融合，可以有效地解决分割过程中信息丢失和分割不准确的问题。

那么，为什么要做这样特征融合呢？因为每一次下采样提取特征的过程中，必然会损失一些边缘特征，而失去的特征并不能从上采样中找回。并且直接对特征图进行上采样，并没有增加特征信息。所以，为了能够补充更多的特征信息，U-Net将前面的中间变量拼接到后面上采样的结果中，使得特征更加丰富。

对于特征融合有两种做法：

第一种是Add操作，类似于ResNet，将两个特征变量进行相加。使得原来的变量包含更多的信息。
第二种是Concat操作，将特征变量在通道的维度上进行拼接，使用torch.cat((x1,x2), dim=1)使得特征信息的增加。U-Net选用的是这种方案。

4. 一些细节

4.1 输出特征图尺寸变小

论文中除了最后的输出层，其余所有卷积层统一为 3x3 的卷积核， padding=0, stride=1，即使用 padding=valid 填充算法，没有padding 所以每次卷积之后 Feature Map 的 H 和 W 都会减2。相反，如果使用 padding=same 填充算法，经过3x3卷积之后输出的特征图尺寸不变，最终上采样输出特征图尺寸与输入图片尺寸一致。但是，padding=same 会引入误差的，而且模型越深层得到的 Feature Map 抽象程度越高，受到 padding 的影响会呈累积效应。

由此可见，经过卷积操作后输出特征图的尺寸会有些许变小（边界信息丢失），这就是为什么 concat Feature Map 的时候需要 crop 的原因。为了保证分割的无缝平铺，需要合理地选择输入图像尺寸，以保证池化操作时能被整除。

4.2 Overlap-tile策略

作者的数据集为 512x512，图像经过镜像padding后裁剪会变成 572 x 572，U-Net的输入尺寸为 572x572，mask输出尺寸为388x388。那么，388x388 如何恢复 512x512 呢？可以通过转置卷积或上采样还原成 512x512，但是U-Net采用了 Overlap-tile策略。

如下图所示，假设要预测黄色的区域，则将蓝色区域输入，因为图片经过模型后尺寸缩小，所以需要大一圈。为了更好预测边缘区域，使用镜像 padding，以获得边缘的周边信息。这样的操作会带来图像重叠问题，即某一图像的周围可能会和另一张图片重叠。计算每个地方重叠次数，最后取平均。

在这里插入图片描述

4.3 镜像padding

图像分割U-Net原理及Pytorch实现

//TODO

如何确定镜像 padding 的尺寸？

一个比较好的策略是通过感受野来确定。因为卷积操作会降低 Feature Map 分辨率，但是我们希望 512x512 的图像的边界点能够保留到最后一层 Feature Map。所以我们需要通过 padding 操作增加图像的分辨率，增加的尺寸即是感受野的大小，也就是说每条边界增加感受野的一半作为镜像 padding。根据图1中所示的压缩路径的网络架构，我们可以计算其感受野：
$\text{rf}=(((0\times2+2+2)\times2+2+2)\times2+2+2)\times2+2+2=60$
这就是为什么U-Net的输入数据是 572x572。

5. U-Net优点

U-Net避免了直接在高级 Feature map 中进行损失计算和监督，而是将低级 Feature map 中的特征结合起来，因此能够保证最终得到的 Feature map 中既包含 high-level 的 Feature，也包含大量的 low-level 的 Feature，最终实现不同尺度下 Feature 的融合，进而提高模型的精确度。

三、相关经验

1. (PyTorch)代码实现

pytorch实现Unet

U-Net具体实现，输入是 Bx1 x 572 x 572，输出是 B x 2 x 388 x 388，输出2个通道代表一个预测前景目标，一个预测背景，然后哪个大，就归为哪一类。当然也可以输出一个通道，然后经过 sigmoid 转成概率，然后大于0.5，就是前景目标。

# sub-parts of the U-Net model
import torch
import torch.nn as nn
import torch.nn.functional as F



# 两个重复的 3x3 的卷积层（`no padding`）
class double_conv(nn.Module):
    '''(conv => BN => ReLU) * 2'''

    def __init__(self, in_ch, out_ch):
        super(double_conv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        x = self.conv(x)
        return x

# 实现左边第一行的卷积
class inconv(nn.Module): 
    def __init__(self, in_ch, out_ch):
        super(inconv, self).__init__()
        self.conv = double_conv(in_ch, out_ch)
    def forward(self, x):
        x = self.conv(x)
        return x

# 下采样
class down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(down, self).__init__()
        self.mpconv = nn.Sequential(
            nn.MaxPool2d(2),
            double_conv(in_ch, out_ch)
        )
    def forward(self, x):
        x = self.mpconv(x)
        return x

# 上采样
class up(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(up, self).__init__()
        if bilinear:
            # 采用双线性插值进行上采样
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            # 采用转置卷积进行上采样
            self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)

        self.conv = double_conv(in_ch, out_ch)
        
    def forward(self, x1, x2):  
        x1 = self.up(x1)

        diffY = x1.size()[2] - x2.size()[2]  # 得到图像x2与x1的H的差值，56-64=-8
        diffX = x1.size()[3] - x2.size()[3]  # 得到图像x2与x1的W差值，56-64=-8

        # 用第一次上采样为例,即当上采样后的结果大小与右边的特征的结果大小不同时，通过填充来使x2的大小与x1相同
        # 对图像进行填充(-4,-4,-4,-4),左右上下都缩小4，所以最后使得64*64变为56*56
        x2 = F.pad(x2, (diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2))

        # for padding issues, see
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd

        # 将最后上采样得到的值x1和左边特征提取的值进行拼接,dim=1即在通道数上进行拼接，由512变为1024
        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x

    
# 实现右边的最高层的最右边的卷积
class outconv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(outconv, self).__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 1)
    def forward(self, x):
        x = self.conv(x)
        return x
    
    
class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        
        self.inc = inconv(in_channels, 64)
        self.down1 = down(64, 128)
        self.down2 = down(128, 256)
        self.down3 = down(256, 512)
        self.down4 = down(512, 1024)
        self.up1 = up(1024, 512)
        self.up2 = up(512, 256)
        self.up3 = up(256, 128)
        self.up4 = up(128, 64)
        self.outc = outconv(64, out_channels)

    def forward(self, x):
        # (1, 572, 572) -> (64, 568, 568)
        x1 = self.inc(x)
        # (64, 568, 568) -> (128, 280, 280)
        x2 = self.down1(x1)
        # (128, 280, 280) -> (256, 136, 136)
        x3 = self.down2(x2)
        # (256, 136, 136) -> (512, 64, 64)
        x4 = self.down3(x3)
        # (512, 64, 64) -> (1024, 28, 28)
        x5 = self.down4(x4)
        # (1024, 28, 28) -> (512, 52, 52)
        x = self.up1(x5, x4)
        # (512, 52, 52) -> (256, 100, 100)
        x = self.up2(x, x3)
        # (256, 100, 100) -> (128, 196, 196)
        x = self.up3(x, x2)
        # (128, 196, 196) -> (64, 388, 388)
        x = self.up4(x, x1)
        # (64, 388, 388) -> (2, 388, 388)
        x = self.outc(x)
        return x
        # return F.sigmoid(x) #进行二分类

网络结构打印输出

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 570, 570]             640
            Conv2d-2         [-1, 64, 568, 568]          36,928
         MaxPool2d-3         [-1, 64, 284, 284]               0
            Conv2d-4        [-1, 128, 282, 282]          73,856
            Conv2d-5        [-1, 128, 280, 280]         147,584
         MaxPool2d-6        [-1, 128, 140, 140]               0
            Conv2d-7        [-1, 256, 138, 138]         295,168
            Conv2d-8        [-1, 256, 136, 136]         590,080
         MaxPool2d-9          [-1, 256, 68, 68]               0
           Conv2d-10          [-1, 512, 66, 66]       1,180,160
           Conv2d-11          [-1, 512, 64, 64]       2,359,808
        MaxPool2d-12          [-1, 512, 32, 32]               0
           Conv2d-13         [-1, 1024, 30, 30]       4,719,616
           Conv2d-14         [-1, 1024, 28, 28]       9,438,208
  ConvTranspose2d-15          [-1, 512, 56, 56]       2,097,664
           Conv2d-16          [-1, 512, 54, 54]       4,719,104
           Conv2d-17          [-1, 512, 52, 52]       2,359,808
  ConvTranspose2d-18        [-1, 256, 104, 104]         524,544
           Conv2d-19        [-1, 256, 102, 102]       1,179,904
           Conv2d-20        [-1, 256, 100, 100]         590,080
  ConvTranspose2d-21        [-1, 128, 200, 200]         131,200
           Conv2d-22        [-1, 128, 198, 198]         295,040
           Conv2d-23        [-1, 128, 196, 196]         147,584
  ConvTranspose2d-24         [-1, 64, 392, 392]          32,832
           Conv2d-25         [-1, 64, 390, 390]          73,792
           Conv2d-26         [-1, 64, 388, 388]          36,928
           Conv2d-27          [-1, 2, 388, 388]             130
================================================================
Total params: 31,030,658
Trainable params: 31,030,658
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.25
Forward/backward pass size (MB): 1096.59
Params size (MB): 118.37
Estimated Total Size (MB): 1216.21
----------------------------------------------------------------
ckward pass size (MB): 1096.59
Params size (MB): 118.37
Estimated Total Size (MB): 1216.21
----------------------------------------------------------------