[Transformer]Transformer-Unet: Raw Image Processing with Unet

Unet)

Abstract

医学图像分割任务在医学图像分析领域十分重要，精确的分割结果可以帮助医生进行诊断，提升患者就诊体验。UNet是最被广泛采用的用于医学图像分割的框架，因为通过其编解码结构结合skip connection可以有效的保留原始特征；而Transformer作为在NLP领域中一种十分流行的模型，目前正逐步引入到计算机视觉领域，在目标检测、图像分类和图像分割任务中均取得了优异的效果。因此将二者结合起来可以比单独使用两种方法更有效。

本文提出的Transformer-UNet网络结构，在UNet中加入Transformer模块，Transformer直接处理原始图像，而不是UNet提取的特征图。并在CT82数据集上测试其对胰腺分割的效果，实验结果表明这种端到端的网络比以往许多机遇UNet的算法搜获得了更好的分割结果。

Section I Introduction

CNN作为计算机视觉领域的主流网络框架，在诸多任务中达到了SOTA。对于分割任务，许多基于CNN的方法超过了传统的基于统计模型（MRFs,CRFs），但是这两类方法都各有千秋，比如CNN擅长通过权重共享和滑窗提取局部特征，统计学方法善于通过建模像素对之间的关系建立全局信息。因此将二者的优势结合起来是十分有用的，也有研究工作证明了这一点。


Transformer最初在NLP领域中被用于机器翻译任务，它可以有效的对整个序列进行建模而不会丢失任何有效信息。目前已有诸多研究尝试将Transformer用于计算机视觉任务，比如将输入图像切分成不同部分组成序列，这类方法天然在捕获全局特征方面表现良好，因为Transformer可以对序列的不同元素之间捕获联系。



但是不管是CNN还是Transformer都有其自身局限性，比如CNN不善于提取全局特征，而transformer计算成本高且不能有效的捕获局部特征，此外数据集也对Transformer有很大影响，小型数据集上不好收敛。如果CNN和Transformer结合则可以互补，从而更适合分割任务。



作为统计模型的替代方法，Transformer可以提高CNN的适用性，因此本文建议在传统的UNet分割网络中加入Transformer层，形成一个端到端的语义分割网络。尽管TransUNet已经在这方面做出了探索，但是TransUNet使用的是feature map，这可能不是最高效的，虽然CNN可以用更大的kernel size来捕获patch之间的关系，但这种方法并没有完全突出变压器之间的优势。



本文认为原始图像更能显示像素之间的关系，对医学图像分析也更为重要。因此本文使用原始图像送入Transformer，然后在搭建的UNet网络中加入transformer层，decoder部分则会将transformer层的结果和CNN的结果进行级联，本文将这种网络称之为TUNet（Transformer-UNet）,在 CT82数据集中超过了UNet,AttenUnet和TransUNet。

Section II Related Work

CNN
目前CNN已经广泛用于计算机视觉任务，如目标检测、人体关键点识别等，。目前需要基于CNN的模型一般层数很深，包含上百万的可训练参数。

Semantic Segmentation

语义分割表示将图像中的每一个像素点进行分类，主要依据全卷积神经网络完成这种密集预测任务。目前主流的语义分割网络有两类，一类是编码器-解码器的网络结构；一类是将CNN与统计模型相结合。


Attention


注意力机制最初在NLP任务中作为RNN的一种替代方案，因为self-attention可以在不丢失信息的情况下捕获序列不同元素之间的关系，而注意力机制也逐渐被引入计算机视觉领域，可以与CNN结合更好的提取信息。
因为CNN与Transformer各有千秋，因此近期有的研究正致力于将二者结合，如TransUnet将Transformer与UNet结合提升了其捕获全局特征的能力，但是很少有研究关注Transformer如何在原始图像汇总发挥作用，因为在原图上使用self-attention计算成本很高；在特征图使用Transformer因为feature map一般尺寸更少，计算量少一些。但这一定程度上限制了神经网络的性能，因为不同patch之间的关系 在原始图像中更清晰，所以我们尝试直接在原始图中使用Transformer，也取得了不错的效果。

Section III Method

本文首先遵循UNet原文搭建基础的CNN结构，其中上采样使用双线性插值，下采样使用最大池化。为了便于插入注意力模块，本文设计了一个几乎对称的网络，但是并没有在encoder和decoder之间设置直接相连，这将在后文进行解释。


模型需要输入的是1D序列，因此原始CHW的图像本文展平为了C * n * n的图像，其中n就是patch size，序列长度为HW/n^2.本文中为了简化设置为H=W正方形。

Fig 1展示了原始Transformer和ViT之间的区别，ViT在MHA和MLP之前进行的layer norm，确保要处理的值不会太大；除此之外的MHA和MLP层与原始Transformer保持一致；本文同样进行了positional embedding。
其他修改：将MLP中的激活函数GELU替换成了ELU，因为在实验中使用ELU效果更好，本文认为ELU有用的原因在于在CT图像中正值和负值一样重要，其中超参数α设置为1.

最终Transformer的输出结果表述为：

xi代表不同的patch，m代表transformer的层数。
对于原始图像，则是在ViT中对整个图像应用1x1卷积然后送入Transformer，依旧遵循Equation 2.




由于Transformer并不擅长捕获局部信息，因为本文在UNet中额外添加了一个encoder，这一encoder不会直接连接到decoder，而是会将其在不同感受野下提取到的feature map，与decoder中的feature map进行级联。
详请参见Fig 2，decoder取Transformer的输出reshape后作为输入，同时在上采样过程中级联encoder的结果，这样确保decoder的输入能够包含不同乳香patch的信息，也可以更好的做预测。

Section IV Implementation

因为本文会直接使用原始图像，那么元树图像的大小、patch size的大小就变得十分重要，决定了transformer模型最终的计算量和运行速度。


数据集：CT82 图像分辨率512 * 512 
patch size = 16
因此序列长度为1024（HW/n^2）
因此decoder的输入是10241616的序列，本文reshape成1512512
权衡了效率和精度


本文将attention headshe设置为8 total layer设置为6
Section V Experiment

本文采用的损失函数是交叉熵损失函数（BCE）



本文的输入图像没有经过任何预处理，因此我们在整个模型中用端到端训练替代了传统的预处理-微调模式；此外TUNet中主要的特征提取器是Transformer模块，因此本文设计的UNet网络其编码器比解码器更浅，一方面这样减少了网络模型一方面尽可能保证decoder部分的性能。这也是我们称之为“准对称”的网络结构。

训练细节：使用自适应调整学习率的Adam优化器,训练120epoch
训练卡：NVIDIA 3090 24GB

数据集：CT82 本文则使用的是CT切片

因为（1）使用切片可以有效扩充数据集数目
;

（2）Transformer的计算会消耗大量计算资源，更适合处理2D图像，因为这样不会显著增加权重文件的大小。

与TUNet作对比的网络有：UNet,Attention UNet,TransUNet

实验结果

分割指标为mIOU和dice系数
CT82数据集划分为 training:testing = 60:22
分割阈值设定为0.8，超过0.8的像素点被分割为胰腺，此外本文实验不仅考虑胰腺分割的精度还会考虑背景的分割精度。

Fig 3展示了使用Transformer做分割的结果，可以看到A是Transformer第一层的输出结果，B是Transformer最后一层的输出结果，Transformer可以在较少层次就捕获全局抽象特征。

Table 1展示了本文与其他网络的分割精度对比，可以看到TUNet的性能超过了基准网络、AttenUNet等。
Table 2则展示了不同模型的参数数量、推理时间对比，TUNet并没有显著增加模型规模和推理速度。

Fig 4展示了不同分割网络的可视化结果。而TUNet借助Transformer，即使在相关像素分隔较远时也能产生良好的分割结果。

Model Variance

本文还测试了使用不同patch size的效果，patch size分为为16,32,64，对比结果参见Table 3.可以看到patch size=16时效果最佳，而更高分辨率可能会降低Transformer的效率，因为patchsize越大，输入序列长度越小，而序列长度对self-attention的计算十分重要。

Table4则展示了使用更浅的基准网络对最终分割结果的影响，可以看到TUNet此时并不具备明显优势，可能是因为Transformer提取的抽象特征需要更深的模型进行编码。

Section V
 Conclusion

本文提出了基于Transformer的UNet模型用于医学图像分割，好处是直接对原始图像进行处理。实验结果表明TUNet性能优于其他基于UNet的网络，但是TUNet应用于较浅层模型时并没有太大的提升UNet的效率，本文预计在未来对这一方面做进一步的改进。