U-Net 是图像分割领域的关键架构之一。本文将探讨 Ronneberger 等人于 2015 年提出的原始 U-Net 架构。我们将学习收缩路径、扩张路径和跳跃连接,以及它们如何组合成 U 形。此外,我们还将讨论现代 U-Net 是如何由各种卷积网络主干构建的。
所以,看完这篇介绍之后,你就会明白:
- 图像分割是在像素级别进行的,用于分类和回归。
- 如何使用跳过连接通过收缩和扩展路径将 U-Net 用于图像分割。
- 目前在类似 U-Net 的架构中使用了哪些主干。
让我们来看看吧!
图像分割概述
经典的计算机视觉方法采用专注于分类的深度学习解决方案。Krizhevsky 等人 (2012) 创建的 AlexNet 就是一个例子。它使用卷积层进行特征学习,然后附加一组紧密连接的层来为输入图像分配类别。
这使得模型能够区分猫和狗、马和牛、汽车和卡车,仅举几个例子。
但是,图像的功能远不止于此!例如,它可以检测图像中的物体。如果您已经熟悉深度学习和计算机视觉,那么您很可能听说过 YOLO 或Transformers等方法,它们可以在图像中绘制物体的方框。
然而,图像的精度应该可以进一步提升。经典的物体检测器绘制边界框,而图像分割方法则在像素级别执行。换句话说,每个像素都会被检查并分配一个类别。这允许你绘制非常清晰的边界框来区分物体。
像素级分类
如果您之前构建过分类模型,那么您会知道 Softmax 用于生成目标类别的伪概率分布。然后,您可以简单地取 Softmax 输出张量的最大参数来找到样本所属的类别。这被称为分割掩码。
在图像分割网络的尾部,输出张量并不代表整幅图像,而是像素级别的输出。这意味着 Softmax 将在像素级别应用,并且你可以为每个像素取最大参数来找到它所属的类。
对于回归也是如此
如果我们简单地省略 Softmax,并将输出通道数设置为 1(或者,如果存在多个因变量,则大于 1),我们将获得每个像素的线性输出。这与构建回归模型时的操作类似,只不过这次是在像素级别。通过使用 MAE 或 MSE 等损失函数,现在您也可以进行图像分割以进行回归分析。
这也意味着,用于分类的图像分割的预训练模型可以适用于回归场景,只需删除 Softmax 激活即可(或者仅更改损失函数,如果实际上在损失函数中应用了 Softmax - 两者都是可能的)。
回归的图像分割示例如下:
- 像素级深度估计
- 地图价值估算
在本文的其余部分,我们将学习 U-Net,这是一个用于图像分割的模型。它由 Ronneberger 等人在 2015 年的一篇关于生物医学图像分割的论文中提出。在最初的研究中,U-Net 用于分类。
U-Net:高层次视角
下图展示了 U-Net。由于该网络由形似 U 的层组构成,其名称的由来也就不足为奇了。U 形的左侧部分称为收缩路径,而右侧部分称为扩张路径。U-Net 结合跳跃连接,对输入图像进行下采样,以了解其显著特征,并通过上采样重建输入(或派生结果,例如分割蒙版)。
现在让我们看一下各个组件,从收缩路径开始。
Ronneberger 等人 (2015) 在 U-Net 研究中,从常规的卷积神经网络入手。每个卷积网络都被他们称为一个收缩网络 (contracting network )。简单来说,这意味着卷积层(可能还有池化层)会缩小特征图(输出)的大小。
例如,如果常规卷积神经网络的输入是 32x32x3 的图像(实际上,可以是 CIFAR-10 样本),那么一个简单的 2x2 卷积层(采用常规步长和 12 个特征图)将产生 30x30x12 的输出。通过将多个卷积层(可能还有池化层)堆叠在一起,可以将原始输入缩小到例如 10x10x64。
回想一下,这使得卷积神经网络能够学习一个层次结构的特征,从更详细的特征(在网络的起始处)到更粗粒度的特征(朝向网络的尾部)。由于网络看起来有点像金字塔(越往下游走,它的宽度就越小),它会收缩,因此被称为收缩网络。
从收缩到上采样
U-Net 使用了收缩网络,但通过上采样网络进行扩展,以便以特定分辨率重建输出。该网络称为扩展路径。
主要思想[..]是通过连续的层来补充通常的收缩网络,其中池化算子被上采样算子所取代。
请注意,在 U-Net 的当代变体中,输出在宽度和高度维度上始终与输入大小相等(例如,32x32x3 → 32x32xC)。然而,在原始 U-Net 中,情况并非如此!
跳过连接的必要性
现在,假设你已经有一个有效的收缩层和上采样层堆栈,因此它看起来非常类似于 U-Net。现在是时候开始训练了,我们不使用预训练的权重,而是使用例如 He init 或 Xavier init 来初始化权重,具体取决于我们使用的激活函数。
在对图像进行前向传递时,它们首先会通过类似 U-Net 模型的收缩部分。然后,从收缩后的输入中,它们会再次进行上采样以找到例如蒙版。在收缩网络末端从收缩(或更概念化地称为汇总)输入进行上采样时,通过上采样找到看起来像原始图像的图像确实非常困难 — — 因为您是在盲目地进行!在网络的这个点上,除了摘要之外,您完全没有关于原始输入是什么样子的信息。如果您将其想象成在许多门之间做出选择,那么您实际上需要走过每一扇门,看看它是否是正确的门,如果不是,则走回去尝试另一扇门。
为了实现定位,收缩路径的高分辨率特征与上采样输出相结合。
因此,稍微引导一下方向会非常有帮助。在 U-Net 中,这是通过跳跃连接实现的,跳跃连接将收缩网络中每一层的输出传递到上采样网络中的相应层。结合当前的上采样状态,跳跃连接用于进一步上采样。通过将输入图像的高层信息(来自收缩网络的末端)与层级的高层信息相结合,预计结果会更好。
所得到的网络适用于各种生物医学分割问题。
事实证明,它的应用范围远不止生物医学分割问题!Ronneberger 等人(2015)的论文已被引用数万次,是基于深度学习的图像分割领域的关键论文之一。U-Net 是一种广泛使用的架构,至今仍是图像分割领域的最佳选择之一。
单独的 U-Net 构建块
现在你已经了解了 U-Net 的高层架构,不妨看看它的各个组件。我们从收缩路径开始,然后是扩张路径,最后是跳跃连接。
承包路径
收缩部分由多个卷积层构建块组成,即所谓的卷积块。
卷积块的数量是可配置的,但在原始论文中是五个。
请注意,除最后一个卷积块外,每个卷积块都允许其输出用作跳过连接,因此在五块设置中我们总共有四个跳过连接。
收缩路径遵循卷积网络的典型架构。它由重复应用两个 3x3 卷积(无填充卷积)组成,每个卷积后跟一个修正线性单元 (ReLU) 和一个步长为 2 的 2x2 最大池化操作,用于下采样。
U-Net 收缩路径中的每个卷积块由两个卷积层组成,每个卷积核大小为 3x3。它们没有填充。回想一下,当使用步长为 1 的 3x3 卷积核时,图像的高度和宽度分别减少 2 个像素。这在显示收缩路径的图像中清晰可见。例如,输入图像在卷积块中从 572x572 像素收缩到 570x570 像素,然后又收缩到 568x568 像素。
每个卷积层后面都有一个 ReLU 激活函数。
在每个卷积块之间执行一个最大池化操作。它使用 2x2 的池化大小和步长 2 进行下采样。这意味着输入图像在宽度和高度上都被切成两半。这可以通过将特征图大小翻倍来补偿。由于起始特征图数量为 64,我们最终得到了 5 倍的倍数,因此 U 形底部有 1024 个特征图。
换句话说,相对而言,U-Net 学习的高分辨率特征较少,而学习的低分辨率特征较多。这使得资源使用更加平衡。
通过收缩路径,输入图像被下采样,并学习到大量特征。这些特征现在可以用于上采样,既可以从收缩路径的输出,也可以从跳过连接中获取。现在让我们来看看这是如何实现的。
扩展路径
现在我们了解了收缩路径,现在该看看扩张路径了。它与收缩路径类似,但一切都是反向进行的。
执行的不是对输入图像进行下采样,而是上采样。因此,扩展路径包含的块数量与收缩路径相同,但这些块是上采样块。
这些上采样块的一般过程如下:
- 上采样在特定层级的特征图上执行,以便有效地将特征图大小向上提升一级。上采样还会将特征图的数量减少 2。
- 跳跃连接被中心裁剪到上采样输入的宽度和高度。注意,跳跃连接的特征图数量与上采样输入相同。
- 两者连接起来,先进行跳跃连接,然后进行上采样输入。特征图的数量现在翻倍了。
- 使用两个具有 2x2 核的卷积层,特征图再次被下采样,特征图的数量减少到原始数量。因此,例如,当连接产生 512x2 = 1024 个特征图时,卷积层再次产生 512 个。
对于上采样,使用核大小为 2x2、步长为 2 的上采样层。卷积层再次使用核大小为 3x3、步长为 1 的卷积层,然后使用 ReLU 激活函数。
扩展路径中的每一步都包含以下操作:对特征图进行上采样,然后进行 2x2 卷积(“上卷积”),将特征通道数减半;与收缩路径中相应裁剪的特征图连接;以及两个 3x3 卷积,每个卷积后跟一个 ReLU。(…) 在最后一层,使用 1x1 卷积将每个 64 分量特征向量映射到所需数量的类别。
这意味着最终你会得到一张带有初始数量特征图的上采样图像。根据你的问题(例如,有两个可能类别的分类),现在使用 1x1 核卷积来减少特征图的数量。宽度和高度保持不变,但特征图数量会减少到,比如说 2 个。
当然,之后可以进行 Softmax 激活(可以作为单独的层,或者最好将其添加到损失函数中)和分类交叉熵损失,以训练图像分割模型进行分类。如果希望使用 U-Net 进行回归,我们可以省略 Softmax,让输出为线性的,并且(很可能)只有一个输出特征图。在回归场景中,MSE 和 MAE 等损失函数将是不错的起点。
跳过连接
现在我们已经了解了收缩路径和扩张路径,是时候更详细地研究 U-Net 的最后一个核心方面了。虽然上一节中已经简要提到过,但了解跳过连接的作用仍然很重要。
回想一下关于收缩路径的部分,除了最后一层,每一层都会生成跳跃连接。实际上,在图中你可以看到跳跃连接从上层四个收缩层中产生,而最后一层(作为 U 形的底部)没有生成跳跃连接。
这些跳跃连接随后会在 U 形另一侧的扩展层中重复使用。但请注意,在 U-Net 的原始设置中,跳跃连接的特征图宽度和高度与上采样输入的特征图宽度和高度不相等(实际上,它们更大)。换句话说,如果要构建 U-Net,则无法直接将它们连接起来。
Ronneberger 等人(2015)在他们的研究中解决了这个问题,他们从跳跃连接中提取一个中心裁剪区域——一个与上采样输入大小相同的裁剪区域。例如,假设跳跃连接的特征图尺寸为 136 x 136 像素,而上采样输入尺寸为 104 x 104 像素,则 136 x 136 像素的图像会通过从上、下、左、右四个方向移除 (136-104)/2 个像素,中心裁剪为 104 x 104 像素。现在,两者可以连接起来了。
由于每次卷积都会丢失边界像素,因此裁剪是必要的。
U-Net 极大地受益于跳跃连接,因为它将“网络在收缩路径的特定层级上看到的”信息传递到扩张路径。回想一下门的场景:如果你要穿过成千上万扇门,你可以盲目地走,但最终会走错很多路——导致你花费很长时间才能找到正确的方向。如果你能有所引导,你可能会犯一些判断错误,但你能更快地找到正确的门。
实用说明:现代 U-Nets 中的不同骨干网络
到目前为止,我们已经了解了 Ronneberger 等人在原始著作中如何实现 U-Net 架构。多年来,许多人使用过不同的 U-Net 设置,包括在 ImageNet 等数据集上进行预训练,然后根据特定的图像分割任务进行微调。
这意味着今天,您可能会使用不再使用上述原始架构的 U-Net - 但它仍然是一个很好的起点,因为收缩路径、扩展路径和跳过连接保持不变。
目前,U-Net 架构的常见主干网络包括 ResNet、ResNeXt、EfficientNet 和 DenseNet。这些网络通常已在 ImageNet 数据集上进行预训练,因此许多常见特征都已学习到。使用这些以预训练权重初始化的主干 U-Net,您的分割问题很可能能够更快地达到收敛。
U-Net 在现代机器学习中的应用
TU-Net 架构在机器学习中仍然具有重要意义,尤其是在图像处理和计算机视觉领域。它最初是为生物医学图像分割而设计的,但由于其在处理需要精确定位的任务方面表现出色,现已应用于各个领域,例如:
1.医学成像:
- U-Net 常用于医学图像分割,尤其适用于肿瘤检测、器官分割和细胞追踪等任务。它能够生成精细的像素级分割,非常适合需要高精度的医疗任务。
2.卫星图像分割:
- U-Net 用于分析卫星图像,以实现土地覆盖分类、检测航空图像中的物体以及监测环境变化等目的。
3.自动驾驶
- 在自动驾驶汽车领域,U-Net 常用于道路分割、车道检测和障碍物识别。其捕捉精细细节的能力对于确保安全至关重要。
4.内容生成与恢复:
- U-Net 还适用于图像到图像的转换任务,例如图像超分辨率、图像去噪和修复。
5.3D图像处理:
- U-Net 的变体已被开发用于处理 3D 数据,这在体积医学成像和 3D 点云处理等领域非常重要。
6.一般细分:
- U-Net 一直是一般分割任务的首选架构之一,因为它能够生成精确的输出掩码并处理大输入尺寸。
为什么 U-Net 仍然具有相关性
- 跳过连接:U-Net 中收缩(编码器)路径和扩展(解码器)路径之间的跳过连接有助于保留下采样过程中可能丢失的空间信息。这种结构有助于维持高质量、细粒度的分割。
- 灵活性:U-Net 非常灵活,可以根据使用情况调整为不同的版本,如 3D U-Net、Attention U-Net、ResUNet 等。
- 性能:即使出现了更新的架构,U-Net 通常仍能提供具有竞争力的性能,尤其是在医疗和基于分割的应用中。
总而言之,U-Net 在机器学习领域仍然具有很高的相关性,并且仍然是分割相关任务的有力选择。其适应性强且久经考验的记录使其成为许多现代应用中的常用模型。
U-Net图像分割实战项目(教程与课件)+ 资料包↓(或看我个人简介处)