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. 方法
给定图像 ,空间分辨率为 H × W,通道数为 C。目标是预测相应的大小为 H × W 的像素级标签图。最常见的方法是直接训练 CNN(例如 UNet),首先将图像编码为高级特征表示,然后将其解码回全空间分辨率。
与现有方法不同,提出的方法通过使用 Transformer 将自注意力机制引入编码器设计中。
3.1 Transformer 作为编码器
图像序列化
我们首先通过将输入 x 重塑为一系列扁平化的 2D 补丁 来执行标记化; 其中每个补丁的大小为 且 是图像补丁的数量(即输入序列长度)。
补丁嵌入
使用可训练的线性投影将矢量化补丁 映射到潜在的 D 维嵌入空间。为了对补丁空间信息进行编码,我们学习特定位置嵌入,将其添加到补丁嵌入中以保留位置信息。
其中 是块嵌入投影, 表示位置嵌入。
Transformer 编码器由 L 层多头自注意力 (MSA) 和多层感知器 (MLP) 块组成。因此第 层的输出可以写成如下:
其中 LN(·) 表示层归一化算子, 是编码图像表示。
3.2 TransUNet
出于分割目的,直观的解决方案是简单地将编码特征表示 上采样到全分辨率,以预测密集输出。为了恢复空间顺序,编码特征的大小首先应从 到 。使用 1 × 1 卷积将重构特征的通道大小减少到类数,然后将特征图直接双线性上采样到全分辨率 H × W 以预测最终的分割结果。
因为 通常远小于原始图像分辨率 ,因此不可避免地会导致低级细节的丢失(例如器官的形状和边界)。为了补偿这种信息损失,TransUNet 采用混合 CNN-Transformer 架构作为编码器以及级联上采样器来实现精确定位。
CNN-Transformer 混合作为编码器
TransUNet 采用 CNN-Transformer 混合模型,其中 CNN 首先用作特征提取器,为输入生成特征图。补丁嵌入应用于从 CNN 提取的特征图而不是原始图像中提取的 1 × 1 补丁。
1) 它允许我们在解码路径中利用中间高分辨率 CNN 特征图;
2)我们发现混合 CNN-Transformer 编码器比简单使用纯 Transformer 作为编码器表现更好。
级联上采样器
级联上采样器(CUP),它由多个上采样步骤组成,用于解码隐藏特征以输出最终的分割掩模。将隐藏特征 的序列重塑为 的形状后,通过级联多个上采样块来实例化 CUP,以达到从 到 的全分辨率,其中每个块依次由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 享有高级全局上下文信息和低级细节的优势。