5.23.9 TransUNet：Transformers 为医学图像分割提供强大的编码器

news2024/7/6 19:37:42

TransUNet，它兼具 Transformers 和 U-Net 的优点，作为医学图像分割的强大替代方案。一方面，Transformer 对来自卷积神经网络 (CNN) 特征图的标记化图像块进行编码，作为用于提取全局上下文的输入序列。另一方面，解码器对编码特征进行上采样，然后将其与高分辨率 CNN 特征图相结合以实现精确定位。

1. 介绍

U-Net 由对称的编码器-解码器网络组成，具有跳跃连接以增强细节保留，已成为事实上的选择。基于这种方法，在广泛的医学应用中取得了巨大的成功，例如磁共振（MR）的心脏分割、计算机断层扫描（CT）的器官分割和息肉从结肠镜检查视频中分割。

与之前基于 CNN 的方法不同，Transformers 不仅在建模全局上下文方面功能强大，而且在大规模预训练下也表现出了对下游任务的卓越可迁移性。

Transformer 将输入视为一维序列，并专门专注于对所有阶段的全局上下文进行建模，因此导致低分辨率特征缺乏详细的本地化信息。并且该信息无法通过直接上采样到全分辨率来有效恢复，因此会导致粗分割结果。另一方面，CNN 架构（例如 U-Net ）提供了一种提取低级视觉线索的途径，可以很好地弥补这种精细的空间细节。

1.1 TransUNet

该框架从序列到序列预测的角度建立了自注意力机制。为了弥补Transformers带来的特征分辨率损失，TransUNet采用了混合CNN-Transformer架构，既利用了来自CNN特征的详细高分辨率空间信息，也利用了Transformers编码的全局上下文。

Transformers编码的自注意力特征随后被上采样，与编码路径中跳过的不同高分辨率CNN特征相结合，以实现精确定位。与以前基于CNN的自注意力方法相比，我们基于Transformer的架构提供了一种更好的利用自注意力的方法。此外，我们观察到，更深入地结合低级特征通常会带来更好的分割精度。

2. 相关工作

2.1 将 CNN 与自注意力机制相结合

尝试通过基于特征图对所有像素的全局交互进行建模，将自注意力机制集成到 CNN 中。基于编码器-解码器 u 形架构，提出了集成到跳跃连接中的附加注意门模块。与这些方法不同的是，我们使用 Transformer 将全局自注意力嵌入到方法中。

2.2 Transformers

Vision Transformer (ViT) 通过直接将具有全局自注意力的 Transformer 应用于全尺寸图像，实现了 ImageNet 分类的最先进技术。

3. 方法

给定图像 $\mathbf{x}\in\mathbb{R}^{H\times W\times C}$ ，空间分辨率为 H × W，通道数为 C。目标是预测相应的大小为 H × W 的像素级标签图。最常见的方法是直接训练 CNN（例如 UNet），首先将图像编码为高级特征表示，然后将其解码回全空间分辨率。

与现有方法不同，提出的方法通过使用 Transformer 将自注意力机制引入编码器设计中。

3.1 Transformer 作为编码器

图像序列化

我们首先通过将输入 x 重塑为一系列扁平化的 2D 补丁 $\mathbf{x}_{p}^{i}\in\mathbb{R}^{P^2\cdot C}|i=1,...,N$ 来执行标记化；其中每个补丁的大小为 $P\times P$ 且 $N=\frac{HW}{P^2}$ 是图像补丁的数量（即输入序列长度）。

补丁嵌入

使用可训练的线性投影将矢量化补丁 $x_p$ 映射到潜在的 D 维嵌入空间。为了对补丁空间信息进行编码，我们学习特定位置嵌入，将其添加到补丁嵌入中以保留位置信息。

$\mathbf{z}_{0}=[\mathbf{x}_{p}^{1}\mathbf{E}; \mathbf{x}_{p}^{2}\mathbf{E};\cdots; \mathbf{x}_{p}^{N}\mathbf{E}]+\mathbf{E}_{pos}$

其中 $\mathbf{E}\in\mathbb{R}^{(P^{2}\cdot C)\times D}$ 是块嵌入投影， $\mathbf{E}_{pos}\in\mathbb{R}^{N\times D}$ 表示位置嵌入。

Transformer 编码器由 L 层多头自注意力 (MSA) 和多层感知器 (MLP) 块组成。因此第 $\ell$ 层的输出可以写成如下：

$\mathbf{z}_{\ell}^{\prime}=\mathrm{MSA}(\mathrm{LN}(\mathbf{z}_{\ell-1}))+\mathbf{z}_{\ell-1}$

$\mathbf{z}_{\ell}=\mathrm{MLP}(\mathrm{LN}(\mathbf{z}_{\ell}^{\prime}))+\mathbf{z}_{\ell}^{\prime}$

其中 LN(·) 表示层归一化算子， $\mathbf{z}_L$ 是编码图像表示。

3.2 TransUNet

出于分割目的，直观的解决方案是简单地将编码特征表示 $\mathbf{z}_{L}\in\mathbb{R}^{\frac{HW}{P^{2}}\times D}$ 上采样到全分辨率，以预测密集输出。为了恢复空间顺序，编码特征的大小首先应从 $\frac{HW}{P^{2}}$ 到 $\frac{H}{P}\times\frac{W}{P}$ 。使用 1 × 1 卷积将重构特征的通道大小减少到类数，然后将特征图直接双线性上采样到全分辨率 H × W 以预测最终的分割结果。

因为 $\frac{H}{P}\times\frac{W}{P}$ 通常远小于原始图像分辨率 $H \times W$ ，因此不可避免地会导致低级细节的丢失（例如器官的形状和边界）。为了补偿这种信息损失，TransUNet 采用混合 CNN-Transformer 架构作为编码器以及级联上采样器来实现精确定位。

CNN-Transformer 混合作为编码器

TransUNet 采用 CNN-Transformer 混合模型，其中 CNN 首先用作特征提取器，为输入生成特征图。补丁嵌入应用于从 CNN 提取的特征图而不是原始图像中提取的 1 × 1 补丁。

1) 它允许我们在解码路径中利用中间高分辨率 CNN 特征图；

2）我们发现混合 CNN-Transformer 编码器比简单使用纯 Transformer 作为编码器表现更好。

级联上采样器

级联上采样器（CUP），它由多个上采样步骤组成，用于解码隐藏特征以输出最终的分割掩模。将隐藏特征 $\mathbf{z}_{L}\in\mathbb{R}^{\frac{HW}{P^{2}}\times D}$ 的序列重塑为 $\frac{H}{P}\times\frac{W}{P}\times D$ 的形状后，通过级联多个上采样块来实例化 CUP，以达到从 $\frac{H}{P}\times\frac{W}{P}$ 到 $H \times W$ 的全分辨率，其中每个块依次由2×上采样算子、3×3卷积层、ReLU层组成。CUP 与混合编码器一起形成了一个 u 形架构，可以通过跳跃连接实现不同分辨率级别的特征聚合。

4. 实验与讨论

4.2 实验细节

输入分辨率和块大小P设置为224×224和16。因此，我们需要在 CUP 中连续级联 4 个 2× 上采样块才能达到全分辨率。模型使用 SGD 优化器进行训练，学习率为 0.01，动量为 0.9，权重衰减为 1e-4。

以逐片方式推断所有 3D 体积，并将预测的 2D 切片堆叠在一起以重建 3D 预测以进行评估。

4.5 可视化

纯基于 CNN 的方法 U-Net 和 AttnUNet 更容易对器官进行过度分割或分割不足（例如，在第二行中，脾脏被 AttnUNet 过度分割，而被 UNet 分割不足）
与其他方法相比，我们的 TransUNet 预测的假阳性更少，这表明 TransUNet 在抑制那些噪声预测方面比其他方法更有利。
与基于 Transformer 的模型进行比较，我们可以观察到，R50-ViT-CUP 在边界和形状方面的预测往往比 TransUNet 的预测更粗糙。
TransUNet 能够进行更精细的分割并保留详细的形状信息。原因是 TransUNet 享有高级全局上下文信息和低级细节的优势。

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.coloradmin.cn/o/1689465.html

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈，一经查实，立即删除！