摘要

U-Net（Convolutional Networks for Biomedical Image Segmentation）是一种用于图像分割的深度学习网络，最初设计用于医学图像分割任务。其核心结构由对称的编码器-解码器组成：编码器通过卷积和池化操作逐步提取图像的抽象特征并降低分辨率，从而捕捉目标的全局语义信息；解码器通过上采样和卷积操作逐步恢复分辨率，并结合编码器提供的低层特征图（通过跳跃连接）重建目标的细节信息，从而实现精确的分割。为了解决深层网络中的细节丢失问题，U-Net 引入了跳跃连接，将编码器的低层特征图与解码器的高层特征图拼接，从而保留细节信息并提升分割精度。针对医学图像标注数据有限的问题，U-Net 采用了弹性形变的数据增强技术，增强了模型在少量数据上的泛化能力。然而，U-Net 也存在一些局限性：如对小目标的检测能力有限，对多尺度目标的适应性不足，以及对数据增强的依赖较强。为了解决这些问题，许多方法被提出，例如引入注意力机制（如 CBAM、SE Block）增强特征表达能力，使用多尺度特征融合（如空洞卷积、金字塔池化）提升对多尺度目标的适应性，以及通过自监督学习减少对标注数据的依赖。这些改不仅提升了U-Net的性能，也进一步扩大了其应用范围。

Abstract

U-Net (Convolutional Networks for Biomedical Image Segmentation) is a deep learning network designed for image segmentation, initially developed for medical image segmentation tasks. Its core structure consists of a symmetric encoder-decoder architecture: the encoder gradually extracts abstract features from the image and reduces resolution through convolutional and pooling operations, capturing the global semantic information of the target; the decoder gradually restores resolution through upsampling and convolutional operations, combining low-level feature maps from the encoder (via skip connections) to reconstruct detailed information of the target, thereby achieving precise segmentation. To address the issue of detail loss in deep networks, U-Net introduces skip connections, which concatenate low-level feature maps from the encoder with high-level feature maps from the decoder, preserving detail information and improving segmentation accuracy. To tackle the limited availability of annotated medical image data, U-Net employs data augmentation techniques such as elastic deformation, enhancing the model’s generalization ability with small datasets. However, U-Net also has some limitations: for example, its ability to detect small targets is limited, its adaptability to multi-scale targets is insufficient, and it heavily relies on data augmentation. To address these issues, many methods have been proposed, such as introducing attention mechanisms (e.g., CBAM, SE Block) to enhance feature representation, using multi-scale feature fusion (e.g., dilated convolution, pyramid pooling) to improve adaptability to multi-scale targets, and leveraging self-supervised learning to reduce dependence on annotated data. These improvements not only enhance the performance of U-Net but also further expand its application scope.

---

文章信息

Title：U-Net: Convolutional Networks for Biomedical Image Segmentation
Author：Olaf Ronneberger, Philipp Fischer, and Thomas Brox
Source：https://arxiv.org/abs/1505.04597

---

研究动机

从2012年Alexnet的提出以来，卷积神经网络已经广泛运用于计算机视觉任务。卷积网络的典型用途是分类任务，其中图像的输出是单个类别标签。然而，在许多视觉任务中，尤其是在生物医学图像处理中，期望的输出应当包括定位，即，假设将类标签分配给每个像素。此外，在生物医学任务中，用于训练的数据很少。所以，本文构建了一种全卷积网络用来分割图像，这种网络需要的训练图像很少，却能产生精确的分割结果。

U-Net网络结构

U-Net网络是一个全卷积网络，是一个编码-解码的结构，有基本的对称性。

U-Net的网络架构如上图所示，U-Net 可以分为三部分：
第一部分是主干特征提取部分，遵循经典的卷积网络架构，是卷积和最大池化的堆叠，利用主干特征提取部分我们可以获得五个初步有效特征层。
第二部分是加强特征提取部分，对主干特征提取部分得到的五个初步有效特征层进行逐步的上采样和拼接融合，得到与第一个初步有效特征层有相同通道数的特征层。
第三部分是预测部分，对第二部分得到的特征层进行卷积操作，对每一个像素点分类，得到图像分割结果图。
下面对具体的卷积层和上采样层进行说明：

1. conv $3\times 3$，ReLU：此结构中所有的卷积都是 $stride=1$，$padding=0$，此结构在主干特征提取部分使用，除连接输入的卷积层用的第一个卷积核通道是64外，其他卷积层的通道数都是输入通道数的2倍（通道数加倍以弥补下采样带来的损失）。
2. copy and crop：对主干特征提取部分得到的前四个初步有效特征层进行裁剪，以便与上采样得来的特征层进行拼接。
3. max pool $2\times 2$:做$2\times 2$的最大池化操作（下采样）， $stride=2$，$padding=0$，池化前后的通道数不变，宽高减半。
4. up-conv $2\times 2$：上采样，可用转置卷积或双线性插值等，转置卷积前后宽高加倍，通道数减半。
5. conv $1\times 1$：卷积核为$1\times 1$的卷积操作，$stride=1$，$padding=0$，用在预测部分，卷积后，通道数变为类别数（包括背景类别）。

U-Net 还是一个编码-解码的结构，编码器就是下采样路径，通过卷积和池化操作提取特征。解码器就是上采样路径，通过上采样和跳跃连接恢复分辨率。

U-Net网络搭建

首先，为了方便搭建网络，可以定义一个标准的卷积块。另外吗，本网络不是严格按中文中的网络设置进行搭建的。在实现中，实际上可以在下采样的卷积操作时设置 padding = 1 ，以保持卷积前后的宽高不变，这样在与上采样的结果进行拼接时就不需要裁剪（文中的上采样后宽高加倍），最终得到的分割结果也是和输入图像的尺寸一样，而不是输入图像中间的一部分。

def _block(in_channels, features, name):
        """
        定义一个标准的卷积块。

        参数:
            in_channels (int): 输入通道数。
            features (int): 输出通道数。
            name (str): 块的名称（用于调试）。

        返回:
            nn.Sequential: 包含两个卷积层、两个批归一化层和两个 ReLU 激活函数的序列模块。
        """
        return nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=features,
                kernel_size=3,
                padding=1,  #不同于原文，设置 padding=1 使卷积前后的宽高保持不变
                bias=False,
            ),
            nn.BatchNorm2d(num_features=features),  # 批归一化层
            nn.ReLU(inplace=True),  # ReLU 激活函数
            nn.Conv2d(
                in_channels=features,
                out_channels=features,
                kernel_size=3,
                padding=1,
                bias=False,
            ),
            nn.BatchNorm2d(num_features=features),  # 批归一化层
            nn.ReLU(inplace=True),  # ReLU 激活函数
        )

标准卷积块的参数信息：
in_channels：是卷积块的输入通道数，int型；
features：是卷积块的输出通道数，int型；
name：是块的名称，方便调试，string型。
本函数返回的是经过两组卷积、Batchnormal、ReLu操作后的结果。
接下来搭建编码器部分：
先说明网络中的参数:
in_channels : 输入图像的通道数（例如，灰度图为 1，RGB 图为 3），int型。
out_channels : 输出图像的通道数（例如，二分类任务为 1，多分类任务为类别数），int型。
init_features : 初始特征通道数（决定网络的宽度），int型。

        features = init_features  #定义编码器的通道数

        # 编码器部分（下采样路径）
        # 第一个编码器块
        self.encoder1 = UNet._block(in_channels, features, name="enc1")  
        # 最大池化层
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  
        # 第二个编码器块
        self.encoder2 = UNet._block(features, features * 2, name="enc2")  
        # 最大池化层
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  
        # 第三个编码器块
        self.encoder3 = UNet._block(features * 2, features * 4, name="enc3")  
        # 最大池化层
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)  
        # 第四个编码器块
        self.encoder4 = UNet._block(features * 4, features * 8, name="enc4")  
        # 最大池化层
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)  

编码器和解码器之间还有一个连接的部分：

# 底部（最深层）
        self.bottleneck = UNet._block(features * 8, features * 16, name="bottleneck")  

下面搭建解码器部分：

        # 上采样层
        self.upconv4 = nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2)  
        # 第四个解码器块
        self.decoder4 = UNet._block(features * 16, features * 8, name="dec4")
        # 上采样层 
        self.upconv3 = nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2)  
        # 第三个解码器块
        self.decoder3 = UNet._block(features * 8, features * 4, name="dec3")  
        # 上采样层
        self.upconv2 = nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2)  
        # 第二个解码器块
        self.decoder2 = UNet._block(features * 4, features * 2, name="dec2")  
        # 上采样层
        self.upconv1 = nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2)  
        # 第一个解码器块
        self.decoder1 = UNet._block(features * 2, features, name="dec1")  

下面是输出层（预测层）：

        # 输出层
        self.conv = nn.Conv2d(in_channels=features, out_channels=out_channels, kernel_size=1)  # 1x1 卷积层

整体网络搭建：

import torch
import torch.nn as nn
import torch.nn.functional as F

class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, init_features=32):
        """
        U-Net 初始化函数。

        参数:
            in_channels (int): 输入图像的通道数（例如，灰度图为 1，RGB 图为 3）。
            out_channels (int): 输出图像的通道数（例如，二分类任务为 1，多分类任务为类别数）。
            init_features (int): 初始特征通道数（决定网络的宽度）。
        """
        super(UNet, self).__init__()

        # 定义编码器的特征通道数
        features = init_features

        # 编码器部分（下采样路径）
        self.encoder1 = UNet._block(in_channels, features, name="enc1")  # 第一个编码器块
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层
        self.encoder2 = UNet._block(features, features * 2, name="enc2")  # 第二个编码器块
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层
        self.encoder3 = UNet._block(features * 2, features * 4, name="enc3")  # 第三个编码器块
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层
        self.encoder4 = UNet._block(features * 4, features * 8, name="enc4")  # 第四个编码器块
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层

        # 底部（最深层）
        self.bottleneck = UNet._block(features * 8, features * 16, name="bottleneck")  # 瓶颈层

        # 解码器部分（上采样路径）
        self.upconv4 = nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2)  # 上采样层
        self.decoder4 = UNet._block(features * 16, features * 8, name="dec4")  # 第四个解码器块
        self.upconv3 = nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2)  # 上采样层
        self.decoder3 = UNet._block(features * 8, features * 4, name="dec3")  # 第三个解码器块
        self.upconv2 = nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2)  # 上采样层
        self.decoder2 = UNet._block(features * 4, features * 2, name="dec2")  # 第二个解码器块
        self.upconv1 = nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2)  # 上采样层
        self.decoder1 = UNet._block(features * 2, features, name="dec1")  # 第一个解码器块

        # 输出层
        self.conv = nn.Conv2d(in_channels=features, out_channels=out_channels, kernel_size=1)  # 1x1 卷积层

    def forward(self, x):
        """
        前向传播函数。

        参数:
            x (torch.Tensor): 输入图像张量，形状为 (batch_size, in_channels, height, width)。

        返回:
            torch.Tensor: 输出分割结果，形状为 (batch_size, out_channels, height, width)。
        """
        # 编码器部分
        enc1 = self.encoder1(x)  # 第一个编码器块
        enc2 = self.encoder2(self.pool1(enc1))  # 第二个编码器块
        enc3 = self.encoder3(self.pool2(enc2))  # 第三个编码器块
        enc4 = self.encoder4(self.pool3(enc3))  # 第四个编码器块

        # 底部（最深层）
        bottleneck = self.bottleneck(self.pool4(enc4))  # 瓶颈层

        # 解码器部分
        dec4 = self.upconv4(bottleneck)  # 上采样
        dec4 = torch.cat((dec4, enc4), dim=1)  # 跳跃连接（拼接编码器的特征图）
        dec4 = self.decoder4(dec4)  # 第四个解码器块

        dec3 = self.upconv3(dec4)  # 上采样
        dec3 = torch.cat((dec3, enc3), dim=1)  # 跳跃连接
        dec3 = self.decoder3(dec3)  # 第三个解码器块

        dec2 = self.upconv2(dec3)  # 上采样
        dec2 = torch.cat((dec2, enc2), dim=1)  # 跳跃连接
        dec2 = self.decoder2(dec2)  # 第二个解码器块

        dec1 = self.upconv1(dec2)  # 上采样
        dec1 = torch.cat((dec1, enc1), dim=1)  # 跳跃连接
        dec1 = self.decoder1(dec1)  # 第一个解码器块

        # 输出层
        return torch.sigmoid(self.conv(dec1))  # 1x1 卷积 + Sigmoid 激活函数

    @staticmethod
    def _block(in_channels, features, name):
        """
        定义一个标准的卷积块。

        参数:
            in_channels (int): 输入通道数。
            features (int): 输出通道数。
            name (str): 块的名称（用于调试）。

        返回:
            nn.Sequential: 包含两个卷积层、两个批归一化层和两个 ReLU 激活函数的序列模块。
        """
        return nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=features,
                kernel_size=3,
                padding=1,
                bias=False,
            ),
            nn.BatchNorm2d(num_features=features),  # 批归一化层
            nn.ReLU(inplace=True),  # ReLU 激活函数
            nn.Conv2d(
                in_channels=features,
                out_channels=features,
                kernel_size=3,
                padding=1,
                bias=False,
            ),
            nn.BatchNorm2d(num_features=features),  # 批归一化层
            nn.ReLU(inplace=True),  # ReLU 激活函数
        )

下面对模型进行实例化，测试是否正常：

model = UNet(in_channels=1, out_channels=1, init_features=32)
print(model)

运行的部分结果如下：

可见输出与输入的宽高尺寸一致，各部分也都正常。

数据增强

作者使用弹性形变来进行数据增强，由于数据集的不足，并且数据集是细胞组织的图像，细胞组织的边界每时每刻都会发生不规则的畸变，所以这种弹性形变非常有必要。弹性形变可以让网络学习更稳定的图像特征。
使用随机位移矢量在粗糙的3*3网格上(random displacement vectors on a coarse 3 by 3 grid)产生平滑形变(smooth deformations)。 位移是从10像素标准偏差的高斯分布中采样的。然后使用双三次插值计算每个像素的位移。在contracting path的末尾采用drop-out 层更进一步增加数据。

损失函数

文中使用的损失函数为：pixel-wise softmax + cross_entropy
softmax函数为：

其中 $K$ 代表类别数量， $x$ 代表像素位置， $a_k(x)$ 表示像素 $x$ 预测为 $K$ 的概率。
交叉熵损失函数：

ι：Ω→(1,2,3…K)代表true label， w 是权重
对于细胞的分割，细胞间的间隙需要设置较大的损失权重，而大片的背景区域需要设置相对较小的损失权重，下图是文中的权重热力图。

文中对损失权重的定义为：

其中，$W_c(x)$是用于平衡类别频率的权重函数，$d_1$是与最近细胞的距离，$d_2$是与第二近的细胞的距离。

转置卷积

转置卷积（Transposed Convolution）在语义分割和对抗神经网络（GAN）中较为常见，其主要作用就是上采样。
转置卷积也是一种卷积操作，但不是卷积的逆运算。
卷积操作直观上理解就是卷积核在输入上进行滑动卷积

等效矩阵：每次一个卷积核在一个位置上的卷积操作可以等效为矩阵的乘法：

输入转换成一个向量，每一个等效矩阵转化为一个列向量，然后拼接在一起形成矩阵。

通过输入向量和卷积核矩阵的相乘获得输出向量。输出的向量经过整形便可得到我们的二维输出特征。
将输入记为 $I$,向量化的卷积向量记为 $C$，输出向量记为 $O$，则有：$I^T\ast C=O^T$

转置卷积就是按照此思想还原出输出的形状，注意，转置卷积不是卷积的逆运算，而只是形状上的相反关系：
$$O^T\ast C^T=I^T$$

转置卷积的操作：

1. 在输入特征图元素间填充s-1行、列0（其中s表示转置卷积的步距）
2. 在输入特征图四周填充k-p-1行、列0（其中k表示转置卷积的kernel_size大小，p为转置卷积的padding，注意这里的padding和卷积操作中有些不同）
3. 将卷积核参数上下、左右翻转
4. 做正常卷积运算（填充0，步距1）
   
   下面是$s=2,p=1,k=3$的计算例子：
   
   转置卷积后特征图的大小计算如下：
   
   其中stride[0]表示高度方向的stride，padding[0]表示高度方向的padding，kernel_size[0]表示高度方向的kernel_size，索引[1]都表示宽度方向上的。通过上面公式可看出padding越大，输出的特征矩阵高、宽越小，你可以理解为正向卷积过程中进行了padding然后得到了特征图，现在使用转置卷积还原到原来高、宽后要把之前的padding减掉。
   pytorch中给出的转置卷积函数信息如下：
   
   输出特征图宽高计算如下：
   
   下面对转置卷积进行实验：

import torch
import torch.nn as nn


def transposed_conv_official():
    feature_map = torch.as_tensor([[1, 0],
                                   [2, 1]], dtype=torch.float32).reshape([1, 1, 2, 2])
    print(feature_map)
    trans_conv = nn.ConvTranspose2d(in_channels=1, out_channels=1,
                                    kernel_size=3, stride=1, bias=False)
    trans_conv.load_state_dict({"weight": torch.as_tensor([[1, 0, 1],
                                                           [0, 1, 1],
                                                           [1, 0, 0]], dtype=torch.float32).reshape([1, 1, 3, 3])})
    print(trans_conv.weight)
    output = trans_conv(feature_map)
    print(output)


def transposed_conv_self():
    feature_map = torch.as_tensor([[0, 0, 0, 0, 0, 0],
                                   [0, 0, 0, 0, 0, 0],
                                   [0, 0, 1, 0, 0, 0],
                                   [0, 0, 2, 1, 0, 0],
                                   [0, 0, 0, 0, 0, 0],
                                   [0, 0, 0, 0, 0, 0]], dtype=torch.float32).reshape([1, 1, 6, 6])
    print(feature_map)
    conv = nn.Conv2d(in_channels=1, out_channels=1,
                     kernel_size=3, stride=1, bias=False)
    conv.load_state_dict({"weight": torch.as_tensor([[0, 0, 1],
                                                     [1, 1, 0],
                                                     [1, 0, 1]], dtype=torch.float32).reshape([1, 1, 3, 3])})
    print(conv.weight)
    output = conv(feature_map)
    print(output)


def main():
    transposed_conv_official()
    print("---------------")
    transposed_conv_self()


if __name__ == '__main__':
    main()

transposed_conv_official函数是使用官方的转置卷积进行计算，transposed_conv_self函数是按照上面讲的步骤自己对输入特征图进行填充并通过卷积得到的结果。
终端输出为：

tensor([[[[1., 0.],
          [2., 1.]]]])
Parameter containing:
tensor([[[[1., 0., 1.],
          [0., 1., 1.],
          [1., 0., 0.]]]], requires_grad=True)
tensor([[[[1., 0., 1., 0.],
          [2., 2., 3., 1.],
          [1., 2., 3., 1.],
          [2., 1., 0., 0.]]]], grad_fn=<SlowConvTranspose2DBackward>)
---------------
tensor([[[[0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0.],
          [0., 0., 1., 0., 0., 0.],
          [0., 0., 2., 1., 0., 0.],
          [0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0.]]]])
Parameter containing:
tensor([[[[0., 0., 1.],
          [1., 1., 0.],
          [1., 0., 1.]]]], requires_grad=True)
tensor([[[[1., 0., 1., 0.],
          [2., 2., 3., 1.],
          [1., 2., 3., 1.],
          [2., 1., 0., 0.]]]], grad_fn=<ThnnConv2DBackward>)

Process finished with exit code 0

根据输出可知，官方的转置卷积函数与通过以上的步骤实现的转置卷积的结果一致。

创新性与不足

创新性：

1. 编码器-解码器结构：编码器通过卷积和池化操作逐步提取特征，降低分辨率。解码器通过上采样和卷积操作逐步恢复分辨率。这种结构能够有效地捕捉图像的全局信息和局部细节。
2. 跳跃连接：在编码器和解码器之间引入了跳跃连接，将低层特征图与高层特征图拼接。保留了低层特征的细节信息，有助于精确分割目标边界。
3. 数据增强策略：训练过程中使用了弹性形变等数据增强技术。提高了模型在少量标注数据上的泛化能力。

不足：

1. 对小目标的检测能力有限：尽管 U-Net 通过跳跃连接保留了低层特征，但其低层特征的语义信息不足，在处理极小目标时仍可能丢失细节信息。
2. 对遮挡目标的处理能力有限： 此模型难以区分重叠目标的边界。
3. 长距离依赖性不足：UNet捕捉全局上下文的能力有限，虽然跳连接保留了部分空间信息，但在处理更大范围内的上下文关系时，传统卷积操作仍然难以有效捕捉远距离的依赖性。

总结

U-NetU-Net 采用对称的编码器-解码器设计，编码器通过卷积和池化操作提取特征，解码器通过上采样和卷积操作恢复分辨率。这种架构允许U-Net进行端到端训练，并有效地从有限的数据集中学习。U-Net 最初用于医学图像分割（如细胞分割、肿瘤检测），但其高效的架构和强大的性能使其迅速扩展到其他领域，包括：卫星图像分析、工业检测、自然图像处理等。尽管存在一些局限性，但通过不断的改进和优化，U-Net 及其变体在图像分割任务中仍然具有广泛的应用前景。