TransUNet 论文阅读

date

Apr 1, 2023

Last edited time

Apr 1, 2023 06:20 AM

status

Published

slug

TransUNet论文阅读

一、基本介绍

问题动机医学图像分割是开发医疗保健系统（尤其是疾病诊断和治疗计划）的必要先决条件。在各种医学图像分割任务中，U 形架构（U-Net）已成为事实上的标准，并取得了巨大的成功。但是，由于卷积运算的固有局部性，U-Net 通常在明确建模远程依赖关系方面显示出局限性。设计用于序列到序列预测的 Transformer，已经成为具有先天性全局自注意力机制的替代体系结构，但由于 low-level 细节不足，可能导致定位能力受到限制。

TransUnet 而 TransUnet，它同时具有 Transformers 和 U-Net 的优点，是医学图像分割的强大替代方案。一方面，Transformer 将来自卷积神经网络（CNN）特征图的标记化图像块编码为提取全局上下文的输入序列。另一方面，解码器对编码的特征进行上采样，然后将其与高分辨率的 CNN 特征图组合以实现精确的定位。作者认为，借助 U-Net 的组合，通过恢复局部的空间信息，可以将 Transformers 用作医学图像分割任务的强大编码器。

二、网络结构

解决思路

卷积神经网络（CNN），尤其是全卷积网络（FCN），已经成为医学图像分割中的主流方法。在不同的变体中，由具有跳跃连接的对称编码-解码结构组成的 U-Net，以增强细节保留，成为最佳的选择。基于这种方法，在广泛的医学应用中，例如磁共振（MR）的心脏分割，CT 图像的器官分割和息肉结肠镜检查视频分割中取得了巨大的成功。

尽管基于卷积神经网络的方法具有出色的表现能力，但是由于卷积运算在建模长距离关系方面存在内部局限性，因此这些体系结构通常表现较弱，特别是对于患者的结构纹理、形状和大小方面表现出较大差异的情况。为了克服这一局限性，提出了基于 CNN 特征建立 self-attention 机制的方法。另一方面，为序列到序列预测而设计的 Transformer 已经出现，作为一种替代体系结构，仅依赖于注意力机制。与以前基于 CNN 的方法不同，Transformer 不仅在建模全局上下文方面功能强大，而且在大规模预训练下对下游任务也显示出卓越的可传递性。

本文探讨了 Transformer 在医学图像分割中的潜力。然而，有趣的是，作者发现单纯地使用 Transformer 对标记化的图像块进行编码，然后直接将隐藏的特征表示上采样为完整分辨率的密集输出，无法产生令人满意的结果。这是因为 Transformer 将输入视为一维序列，并且只专注于在所有阶段建模全局上下文，因此会产生缺乏详细定位信息的低分辨率特征。而且，这种信息不能通过直接上采样到全分辨率来有效地恢复，从而导致粗略的分割结果。另一方面，CNN 架构（例如 U-Net）提供了提取低层次视觉线索的途径，可以很好地弥补这种精细的空间细节。

主要方法

为此，作者提出了一个医学图像分割框架 TransUNet。该框架从序列到序列的预测角度建立了 self-attention 机制。为了弥补 Transformer 带来的特征分辨率的损失，TransUNet 采用了 CNN-Transformer 混合结构，以利用来自 CNN 特征的详细高分辨率空间信息以及 Transformer 编码的全局上下文。受 U-Net 设计的启发，框架对 Transformer 编码的自注意特征进行上采样，以此与从编码路径中跳过的不同高分辨率 CNN 特征结合，实现精确的定位。这种设计保留了 Transformer 的优势，并有益于医学图像分割。实证结果表明，相比以前基于 CNN 的自注意方法，基于 Transformer 的体系结构提供了更好的利用自注意的方法。大量实验证明了 TransUNet 在各种医学图像分割任务上相对于其他竞争方法的优越性。

结构详解

给定图像，其中空间分辨率为，通道数为。目标是预测相应像素大小为的标签图。最常见的方法是直接训练 CNN（例如 U-Net），首先将图像编码为高级特征表示，然后将其解码回完整的空间分辨率。与现有方法不同，本文作者使用 Transformer，将自注意力机制引入编码器设计。

为了分割的目的，一个直观的解决方案是简单地将编码的特征表示上采样到全分辨率，以预测密集输出。为了恢复空间顺序，应该首先将编码特征的长度从变为。使用卷积将重新形成的特征的通道尺寸减少到类的数量，然后将特征图直接上采样到全分辨率，用于预测最终的分割结果。

虽然将 Transformer 与传统上采样相结合已经产生了合理的性能，但如上所述，这种策略并不是 Transformer 在分割中的最佳使用，因为通常比原始图像分辨率小得多，因此不可避免地导致低分辨率细节的丢失（例如器官的形状和边界）。

因此，为了弥补这种信息损失，TransUnet 采用了一种混合的 CNN-Transformer 架构作为编码器，并使用级联上采样来实现精确的定位。其中 CNN 首先用作特征提取器，为输入生成特征图。Patch 嵌入应用于从 CNN 特征图而不是从原始图像中提取。选择这种设计是因为：首先，它允许在解码路径中利用中间高分辨率 CNN 特征图；其次，作者发现混合 CNN-Transformer 编码器比简单地使用 Transformer 作为编码器表现更好。

级联上采样（CUP）

该模型由多个上采样步骤组成，用于解码隐藏特征以输出最终的分割 Mask。对隐藏特征的序列进行整形，变成，然后通过级联多个上采样块来实现 CUP，从扩展到全分辨率的。每个块依次包括 2X 上采样、一个 3×3 卷积层和一个 ReLU 层。可以看到，CUP 与混合编码器一起形成了一个 U 形架构，它通过跳跃连接支持不同分辨率级别的功能聚合。

主要结果

①验证 CUP 解码器的有效性，使用 ViT 作为编码器，并分别使用 None 和 CUP 作为解码器比较结果;

②为了证明混合编码器设计的有效性，使用 CUP 作为解码器，并分别比较使用 ViT 和 R50-ViT 作为编码器的结果。

③为了 R50-ViT-CUP 和 TransUNet 进行比较，也用 ImageNet 预训练的 ResNet-50 替换了 U-Net 和 AttnUNet 的原始编码器。

消融实验

💡

The input resolution and patch size P are set as 224×224 and 16, unless otherwise specified.