原文地址：https://arxiv.org/pdf/2102.04306.pdf

收录：CVPR 2021

代码:https://github.com/Beckschen/TransUNet

---

首篇将transformers用于医疗分割的文章

设计的Idea：UNet transformers具体模块：ViT ResNet50 skip connection

个人注释了关键公式和结构图，供大家参考。

---

目录

摘要

1 Introduction

2 related work

2.1 CNN-based self-attention methods

2.2 Transformers

3 method

3.1Transformer as Encoder

3.2TransUNet

4 Experiments and Discussion

4.1 Dataset and Evaluation

4.2 Implementation Details

4.3 Comparison with State-of-the-arts

4.4 Analytical Study

4.5 Visualizations

4.6 Generalization to Other Dataset

---

摘要

在医学图像中， U-Net, 已经成为了 de-facto standard 并取得了巨大的成功。但它有缺陷：explicitly modeling long-range dependency 长期依赖显式建模.Transformers 自然有全局self attention机制，但是因为不充足的low-level特征导致局部定位能力有限。

本文提出TransUNet，以结合Unet和transformers的优点。一方面，为了提取全局contexts，transformers encodes 来自CNN feature map标记化图像patches（tokenized image patches）。另一方面，decoder对encoded取样特征，然后与高分辨率相结合CNN feature maps结合，增加局部精度。

TransUNet 在多器官和心脏分割数据集上取得了巨大的成功multi-organ segmentation and cardiac segmentation。

---

1 Introduction

我们发现一个有趣的事，如果直接将transformer编码后的输出，直接上采样到原分辨率，这样不能达到满意的分割效果。这是因为transformer专注于各个阶段的全局context进行建模，因此产生缺乏细节的局部信息low -resolution features。然后这种特征不能直接通过上采样，被有效地恢复到原分辨率会导致分割结果粗糙。 另一方面, CNN提取低级可视化线索，可以补救这些细微的空间细节。

实验证明，Transformer-based architecture 能更好地利用 self-attention (SA) 相比之前的CNN-based SA methods。此外，我们观察到更密集的地方low-level结合特征，可以达到更高的分割准确率。

---

2 related work

2.1 CNN-based self-attention methods

许多文章都是基于特征的maps，通过对所有像素点的全球交互建模，试图进行SA机制整合到CNN中。Wang et al.非局部设计operator，可插入多个中间卷积层。Schlemper et al.基于编码器-解码器 U 提出形状结构 additive attention gate modules ，其被整合到skip-connections中。与这些方法不同，我们直接采用Transformer整合全局SA在我们的方法中。

2.2 Transformers

在许多NLP在任务中，首先通过[14]提出Transformer，此后在许多NLP在任务中建立了最先进的方法。 为了使Transformer也适用于计算机视觉任务，很多modifications已发表的工作。Parmar et al. [11]只有每个局部邻域query像素上应用SA，而非全局应用。最近，Vision Transformer (ViT) [4]达到state-of-the-art在ImageNet 在分类任务中，通过直接应用 Transformers with global self-attention到 full-sized images. 所以我们提出了TransUnet是第一个基础Transformer医学图像分割框架，其取得巨大成功的ViT上构建。

---

3 method

给定图像，其空间分辨率为H×W，通道数为C。 我们的目标是预测相应的像素级label map，尺寸为H×W。最常见的方法是直接训练CNN（例如UNet），首先将图像encode为high-level feature representations，然后将其decode回full spatial resolution。

与现有方法不同，我们的方法是通过使用的Transformer将SA引入编码器设计。 我们将首先在第三位.1节介绍了如何直接应用Transformer分解后的图像patches特征表示编码。 然后，将在3.详细说明2节TransUNet总体框架。

3.1Transformer as Encoder

• Image Sequentialization.图像序列化。 根据[4]（ViT），首先，我们将输入x重塑为flattened 2D patches ，执行标记化（tokenization）。每个patch大小为P^2，patches的数量为。 相当于把H×W原图转化为NP×P大小的2D图片，C还是不变。
• Patch Embedding 我们使用训练有素的线性投影将矢量化patches 映射到潜在的D维embedding空间。 为了对patches 编码空间信息，我们学习了特定的位置embedding，并加入到patchembedding保留位置信息。公式如下：

where 是patch线性投影，代表位置embedding。

理解：一共有N个，其大小为,  分别都×上E，以转化为维度D的向量。如此以来，得到N×D的向量，再与相加。

Transformer编码器由L层的 Multihead Self-Attention(MSA)和多层感知器(MLP) blocks 组成。因此，L层的输出如下：

 

 Where 表示 layer normalization operator， 是在第L层的编码的图像表示。

理解：首先将上一层的输出，喂入LN，然后再进入MSA，再加上上一层的输出（残差结构），得到。然后做同样的事，不过是进入MLP，得到当前层的输出

Transformer层的结构如图1(a)所示。可见经过linear projection之后，就进入了12个transformer层，每层如左边所示。

图1（a）

3.2 TransUNet

为了分割的目的，一个直观的解决方案是简单地将编码的特征表示即N×D，上采样到全分辨率，以预测密集输出。 在这里，为了恢复空间顺序，编码的特征应该首先重塑大小，从

我们使用1×1的卷积将通道数减小到num of class。然后将特征图直接双线性插值（bilinearly upsampled）到全分辨率H×W。在4.3中我们我们在解码器设计中将此 naive的上采样基线表示为"None"

虽然采用双线性插值已经取得了不错的效果，因为通常比原始图像分辨率H × W 小得多，因此不可避免地导致low level细节的丢失（例如器官的形状和边界）。 因此，为了补偿这种信息损失，TransUnet采用了一种Hybrid CNN-Transformer结构作为编码器以及一个级联的 cascaded upsampler，以实现精确局部化。 proposed TransUNet的概述如图1所示。

图1

• CNN-Transformer Hybrid as Encoder. CNN首先用作特征提取器，生成feature map作为input。Patch embedding 被应用在1×1的Patches上。这些Patches from CNN的feature map，而不是直接从原图得到。

这样设计的原因：1) 在解码path中，可以利用intermediate的高分辨率CNN feature map。2) 这比简单地用pure  Transformer 作为decoder好（这不是废话嘛？）

• Cascaded Upsampler (CUP).  包含了许多下采样步骤去解码hidden特征，以输出最后的分割mask. 在reshape hidden特征到后，（即3.2开头所述的部分），我们示例CUP通过级联多个上采样blocks去达到全分辨率从。每个block包含两个上采样操作，一个3×3卷积层，一个relu。

我们看到CUP 和hybrid encoder组成了一个U型结构，通过skip-connection, 使得不同分辨率水平的特征聚合。

---

结构梳理：

TransUnet的结构与Unet类似，是由一个encoder和一个decoder组成的U型结构。

Encoder部分一共做了4次下采样，缩小到16倍，并加入了Transformer机制，最终得到了N个一维向量。 注意transformer之前做了embbeding，将三维作为二维。, 文章设置的P=16，则N=14*14=196

 

transformer中MLP：经过线性层从768放大到3976，然后relu+dropout，再通过一个线性层缩小回768，再drop out

Decoder部分做了4次上采样，最终将此一维向量恢复成了原来的图像。Encoder和Decoder部分还做了三次跳跃连接。

 这两个图参考两位大佬：

TransUnet: 结构解析_ripple970227的博客-CSDN博客_transunet详解TransUNet_不秃头不成功的博客-CSDN博客_transunetTransUnet: 结构解析_ripple970227的博客-CSDN博客_transunet详解

4 Experiments and Discussion

4.1 Dataset and Evaluation

Synapse multi-organ segmentation dataset1  in the MICCAI 2015 Multi-Atlas Abdomen Labeling Challenge  指标：average DSC (Dice similarity coefficient） 和 average Hausdorff Distance (HD)

Automated cardiac diagnosis challenge 指标：average  DSC

4.2 Implementation Details

for TransUNet, 有两种encoder设计，一种是简单地采用 ViT，一种是结合ResNet-50 和 ViT的,也就是文章一直讲的。所有的transformer的backbone都在Imagenet pretrain了。

4.3 Comparison with State-of-the-arts

对比了四个 previous state-of-the-arts: 1) V-Net [9]; 2) DARR [5]; 3) U-Net [12] and 4) AttnUNet

R50-ViT-CUP加了级联上采样机制CUP后，效果增加不少。

最后的TransUnet是在R50-ViT-CUP基础上，再增加了skip-connections

4.4 Analytical Study

在分析超参数

• The Number of Skip-connections.

• On the Influence of Input Resolution.

如果不降低数据的分辨率，其实效果更好，但是为了实验效率，所以采用224

• On the Influence of Patch Size/Sequence Length

 较高的分割性能通常通过较小的patch大小获得。因为这样得到的N sequence比较长，信息没有被过度压缩

• Model Scaling. 模型规模

base: the hidden size D, number of layers, MLP size, and number of heads are set to be 12, 768, 3072, and 12

larger: 24, 1024, 4096, and 16

可以看出差别不大

4.5 Visualizations

4.6 Generalization to Other Datasets

在心脏数据集ACDC上做了实验，已证明方法的通用性，可见得到的结果和多器官分割差不多。