原文地址:https://arxiv.org/pdf/2102.04306.pdf
收录:CVPR 2021
代码:https://github.com/Beckschen/TransUNet
将用于的文章
具体模块:
摘要
1 Introduction
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中。与这些方法不同,我们在我们的方法中。
2.2 Transformers
在许多NLP在任务中,首先通过[14]提出Transformer,此后在许多NLP在任务中建立了最先进的方法。 为了使Transformer也适用于计算机视觉任务,很多modifications已发表的工作。Parmar et al. [11]只有每个局部邻域query像素上应用SA,而非全局应用。最近,Vision Transformer () [4]达到state-of-the-art在ImageNet 在分类任务中,通过直接应用 Transformers with global self-attention到 full-sized images. 所以我们提出了TransUnet是第一个基础Transformer医学图像分割框架,其。
3 method
给定图像,其空间分辨率为H×W,通道数为C。 我们的目标是。最常见的方法是直接训练CNN(例如UNet),首先将图像encode为high-level feature representations,然后将其decode回full spatial resolution。
与现有方法不同,我们的方法是通过使用的Transformer将SA引入编码器设计。 我们将首先在第三位.1节介绍了如何直接应用Transformer分解后的图像patches特征表示编码。 然后,将在3.详细说明2节TransUNet总体框架。
3.1Transformer as Encoder
- 首先,我们将输入x重塑为flattened 2D patches ,执行标记化(tokenization)。每个patch大小为P^2,patches的数量为。 相当于把H×W原图转化为NP×P大小的2D图片,C还是不变。
- 我们使用训练有素的线性投影将矢量化patches 映射到潜在的D维embedding空间。 为了对patches 编码空间信息,我们学习了特定的位置embedding,并加入到patchembedding保留位置信息。公式如下:
where 是patch线性投影,代表位置embedding。
Transformer编码器由L层的 Multihead Self-Attention(MSA)和多层感知器(MLP) blocks 组成。因此,L层的输出如下:
Where 表示 layer normalization operator, 是在第L层的编码的图像表示。
Transformer层的结构如图1(a)所示。
图1(a)
3.2 TransUNet
为了分割的目的,一个直观的解决方案是简单地将编码的特征表示即N×D,上采样到全分辨率,以预测密集输出。 在这里,为了恢复空间顺序,编码的特征应该首先重塑大小,从
我们使用1×1的卷积将通道数减小到num of class。然后将特征图直接到全分辨率H×W。在4.3中我们我们在解码器设计中将此的上采样基线
虽然采用已经取得了不错的效果,因为通常比原始图像分辨率H × W 小得多,因此不可避免地导致low level细节的丢失(例如器官的形状和边界)。 因此,为了补偿这种信息损失,TransUnet采用了一种结构作为编码器以及一个级联的,以实现精确局部化。 proposed TransUNet的概述如图1所示。
图1
- CNN首先用作特征提取器,生成feature map作为input。Patch embedding 被应用在1×1的Patches上。这些Patches from CNN的feature map,而不是直接从原图得到。
这样设计的原因:1) 在解码path中,可以利用intermediate的高分辨率CNN feature map。2) 这比简单地用pure Transformer 作为decoder好(这不是废话嘛?)
- 包含了许多下采样步骤去解码hidden特征,以输出最后的分割mask. 在reshape hidden特征到后,(即3.2开头所述的部分),我们示例CUP通过级联多个上采样blocks去达到全分辨率从。每个block包含两个上采样操作,一个3×3卷积层,一个relu。
我们看到CUP 和hybrid encoder组成了一个U型结构,通过skip-connection, 使得不同分辨率水平的特征聚合。
TransUnet的结构与Unet类似,是由组成的U型结构。
一共做了4次下采样,缩小到16倍,并加入了Transformer机制,最终得到了N个一维向量。 注意transformer之前做了,将三维作为二维。, 文章设置的P=16,则N=14*14=196
经过线性层从768放大到3976,然后relu+dropout,再通过一个线性层缩小回768,再drop out
做了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
在分析超参数
如果不降低数据的分辨率,其实效果更好,但是为了实验效率,所以采用224
较高的分割性能通常通过较小的patch大小获得。因为这样得到的N sequence比较长,信息没有被过度压缩
are set to be 12, 768, 3072, and 12
24, 1024, 4096, and 16
可以看出差别不大
4.5 Visualizations
4.6 Generalization to Other Datasets
在心脏数据集ACDC上做了实验,已证明方法的通用性,可见得到的结果和多器官分割差不多。