U-Net
U-Net 是一种 U 形编码器—解码器 CNN,其跳跃连接只需少量训练图像即可实现像素级精准分割。内容涵盖架构、变体和代码。
U-Net 是一种用于语义分割的卷积神经网络 (CNN),语义分割的任务是为图像中的每个像素分配一个类别。弗赖堡大学的 Olaf Ronneberger、Philipp Fischer 和 Thomas Brox 于 2015 年在《U-Net: Convolutional Networks for Biomedical Image Segmentation》中提出了 U-Net。编码器会缩小图像尺寸以提取图像内容,解码器会将结果扩展回原始分辨率,而跳跃连接会在每个尺度上将编码器的精细细节传递给解码器,因此输出的目标边界清晰。U-Net 仅用几十张带标注的显微镜图像就能有效训练,在 ISBI 2015 细胞跟踪挑战赛的两个类别中大幅领先,并且至今仍是医学图像分析中的标准模型。此后,它的设计也成为许多扩散模型中的去噪网络。
U-Net 的工作原理#
该网络由两个对称部分组成,形状如同字母 U 的两条臂:
- 收缩路径(编码器):重复使用成对的 3×3 卷积和 ReLU 激活,每对卷积之后接一个 2×2 最大池化层,将分辨率减半。通道数在每个层级翻倍,从 64 增至 1,024;因此,网络以空间细节换取更大的感受野和更丰富的特征。
- 扩展路径(解码器):在每个层级,2×2 上卷积会将分辨率加倍并将通道数减半;随后将结果与编码器同一层级的特征图拼接,再通过两个 3×3 卷积将它们融合。最后,使用一个 1×1 卷积将每个像素的特征映射为类别分数。
原始网络有 23 个卷积层和四次池化操作。它使用不填充的卷积,因此 572×572 的输入图块会生成 388×388 的分割图,而大图像则会被切分成互相重叠的图块进行分割。如今,大多数实现都会对卷积进行填充,使输出尺寸与输入尺寸相同。
跳跃连接#
池化会准确丢弃特征所在的位置信息。到了瓶颈层,网络知道某个区域包含细胞膜,却不知道细胞膜经过哪些像素。跳跃连接会将编码器每个层级的高分辨率特征图跨越 U 形结构复制到解码器,让解码器结合上采样路径提供的“内容”和编码器提供的“位置”。U-Net 会沿通道维度拼接两者,而 ResNet 的残差连接则在一个模块内相加特征。下面的 PyTorch 代码会构建一个编码器层、一个瓶颈层和一个解码器层:
import torch
from torch import nn
# One encoder level, a bottleneck, and one decoder level of a U-Net
encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU())
bottleneck = nn.Sequential(nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
upsample = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
decoder = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())
head = nn.Conv2d(16, 2, kernel_size=1) # 2 classes, e.g. cell and background
image = torch.randn(1, 3, 128, 128)
skip = encoder(image) # full-resolution features kept for the decoder
context = bottleneck(skip) # half resolution, twice the channels, wider context
# Skip connection: concatenate upsampled context with the saved high-resolution features
merged = torch.cat([upsample(context), skip], dim=1)
logits = head(decoder(merged))
print(f"context {tuple(context.shape)} -> merged {tuple(merged.shape)} -> logits {tuple(logits.shape)}")它会输出 context (1, 32, 64, 64) -> merged (1, 32, 128, 128) -> logits (1, 2, 128, 128):16 个上采样后的上下文通道与 16 个跳跃连接通道拼接为 32 个通道,输出则为输入中 128×128 个像素各自生成两个类别分数。
用少量数据训练#
最初用于电子显微镜神经结构的训练集包含 30 张 512×512 像素的图像。U-Net 通过大量数据增强弥补了数据不足,尤其使用了模拟组织弯曲方式的随机弹性形变;它还使用加权交叉熵损失,强调彼此接触的细胞之间狭窄的背景间隙,使相邻细胞最终成为彼此分离的区域。如今,医学分割模型通常使用 Dice 损失,该损失可直接衡量掩码重叠程度,并能处理只占少数像素的小型结构。
变体与衍生模型#
- 3D U-Net 和 V-Net:将 2D 运算替换为 3D 运算,以分割 CT 和 MRI 扫描等体积数据。
- Attention U-Net 和 UNet++:添加注意力门控,或嵌套且密集连接的跳跃路径,以优化传递到解码器的编码器特征。
- nnU-Net:一个可自动为新数据集配置预处理、架构和训练流程的框架。MICCAI 2024 的一项研究发现,在严格验证的条件下,该框架中经过良好调优的 CNN U-Net 在 3D 医学分割任务上仍优于更新的 Transformer 和 Mamba 架构。
- 扩散 U-Net:DDPM 和 Stable Diffusion 1.x 使用 U-Net 预测图像或潜变量中的噪声,并添加注意力层以便提示词引导生成;Stable Diffusion 1.x 的 U-Net 约有 8.6 亿个参数,而 ControlNet 则连接到这个 U-Net 的编码器。更新的生成模型改用扩散 Transformer取代 U-Net。
U-Net 与相关模型的比较#
- 全卷积网络 (FCN):2014 年提出的 FCN 将分类器改造成逐像素预测器,并通过相加类别分数图融合粗略和精细预测。U-Net 则构建了完整的解码器,并在每个层级拼接特征。
- 自编码器:同样采用编码器和解码器结构,但训练目标是让输入经过狭窄的瓶颈层后重建出来。跳跃连接会让自编码器绕过瓶颈层,而 U-Net 需要跳跃连接来生成对齐的逐像素标签。
- 实例分割模型:U-Net 会为每个像素输出一个类别,因此两个接触且属于同一类别的目标会共享一个标签,除非训练或后处理过程将它们分开。
Ultralytics YOLO26 语义分割模型的后缀为 -sem,采用了精简版 U-Net 解码器。它们的颈部会将 1/32 尺度的主干网络输出上采样两次,期间分别在 1/16 和 1/8 尺度拼接主干特征,然后在 1/8 尺度进行像素分类,最后将类别分数上采样 8 倍至输入尺寸。










