[Transformer]Transclaw U-Net: Claw UNet with Transformers for Medical Image Segmentation

Abstract

近年来，计算机辅助诊断已经成为热门的研究话题，基于深度神经网络的方法也在医学图像分割和分类任务中取得了优异的性能，但是受限于卷积操作无法不会长城空间联系，本文提出TransClaw UNet结构，在encoding部分将卷积操作与Transformer操作结合在了一起，其中卷积部分负责提取浅层的空间特征帮助恢复图像上采样后的分辨率；Transformer部分负责对patch进行编码，使用自注意力操作来获取序列之间的全局信息。deociding部分保留了传统的从底部向上采样的结构在多器官分割任务中展示了本文的TransClaw UNet比其他网络结构更为优异的分割性能。

Section I Introduction

随着深度学习的不断发展，计算机视觉以广泛用于制造、自动驾驶、计算机辅助诊断、农业等诸多领域，也引起了国内外学者的广泛关注。许多基于CNN的方法已经用于图像分类和分割任务，能够进行精确的医学图像分割任务用于临床诊断。在诸多网络结构中，UNet的U型结构在医学图像分割领域取得了巨大成功，UNet的encoder部分提供高级语义特征，decoder提供低级的更加细节的特征映射，两阶段通过skip connection进行结合，许多基于UNet的改进也取得了优异的性能，如UNet++,ResNet34-UNet,Channel-UNet，Attention UNet,R2UNet等。充分证明了卷积神经网络的高效性和优越性。
虽然大多数医学图像分割任务都是基于CNN的，但是在CAD中仍然难以满足准确分割的要求，主要是CNN是段导弹的，缺乏对长程关系的建模，使得许多器官的分割效果不是很好。因为每个卷积核只关注自身的特征信息和边界，缺乏在较大范围内的特征融合，为了克服这一问题，许多学者也进行了改进，比如使用自注意力机制、空洞卷积、图像金字塔等，但依然没有考虑到医学图像的长程依赖性。
在自然语言处理中Transformer模型，如BERT,GPT具有长程特征学习能力，Transformer可以有效将分布式卷积操作与多头注意力相结合，从而找到输入之间的关联；许多研究正致力于将Transformer用于CV领域，如ViT,Swin Transformer等，Cao等人提出一种类似UNet的纯Transforemer结构可以通过移窗完成特征提取，在通过patch-expanding完成图像分辨率的恢复。

本文提出了TransClaw UNet将Claw UNet与Transformer结合，可以同时进行细节分割和长程关系的学习，其中卷积操作用来提取特征图，然后对高分辨率的特征图谱进行编码；解码部分保留了Claw UNet的结构，并将深度特征图与浅层特征图相结合用以实现精确定位，最后将对应的解码部分、上采样部分和编码部分组合起来共同用于分辨率的恢复。在腹部CT数据集上的获得了更精确的分割结果，也验证了输入更大的输入图像、更小的patch大小和使用更多的skip connection可以提高模型的性能，但是需要消耗更多的计算资源。

Section II Related Works

CNN及其变体：

作为深度学习的基础框架逐渐取代了传统的机器学习方法。2015年UNet提出后在医学图像分割任务中展现了其优越性，也有越来越多的UNet变体结构被提出，如ResNet34-UNet，channel UNet,UNet++，Attention UNet等，也有的对卷积进行改进，如是用可变形卷积和深度卷积。


自注意力及Transformer:


近年来对如何将自注意力机制或Transformer与CNN相结合也有很多研究，因为自注意力可以弥补CNN长程依赖确实的问题，比如Wang等人在卷积层中插入一些non-local的操作，还有的是在UNet中加入全局注意力机制。


Transformers:


将Transformer与Unet结合存在以下问题，一般UNet的网络结构可以在小规模的数据集上进行训练，但是Transformer却需要大量的数据进行训练，DeiT提出了集中将ViT用于小规模数据集训练的训练策略，本文则是采用Claw UNet作为骨干框架，首先在encoder部分进行卷积运算，再对得到的特征映射送入Transformer处理，对齐进行编码，使用多头注意力机制提取响应的特征，解码部分则将底部上采样特征图与编码后的特征图相结合，用于局部特征的恢复。

Section III Method

本节将介绍Trans Claw UNet的整体结构，包含在encoder部分引入的Transformer结构，从而利用自注意力机制更高的进行全局特征的学习；在像素恢复过程中还将编码部分、上采样和解码部分的结果相结合，从而实现更好的分割结果。

Part 1 Overall Structure

Fig 1展示了TransClaw UNet的整体结构，主要包含 编码部分、自底向上采样部分、解码部分，并且通过skip connection连接起来，并且将Transformer与Convolution在编码部分结合起来。





**encoding part:**

设Iin为输入图像，Ifea为卷积操作后的特征图谱，Ideep为encoding部分最深的图像，则卷积部分的操作可以描述为：conv->BN->ReLU
并且在进行下一次卷积操作之前会进行最大池化来减少分辨率。

上述操作旨在保留低级的特征映射，从而使得医学图像中的浅层信息和轮廓信息可以有效的与解码部分相结合。



为了使得网络可以更好的学习长程依赖关系，本文在卷积之后引入了Transformer机制，表述为：


此时输出会变成H/P * W / P大小，P就是patch size。



**decoding part:**




在解码部分，特征图谱则会经过多次上采样后映射回原始分辨率。
并且本文通过skip connection将编码部分、解码部分和上采样部分的结果相结合，表示为：

每一个上采样操作就会经过3x3卷积将分辨率恢复到H / W * W / 2之后，只有最后一次上采样不需要，如此操作后，图像恢复至原始大小。

Part 2 Transformer in Trans Claw UNet

为了保证图像的原始特征和轮廓，在送入Transformer之前进行了多次卷积，并将提取到的特征图作为Transformer的输入。将输入图像划分为N个PXP大小的patch，他们之间互不重叠，得到图像序列，然后将其reshape到N X P ^ 2的大小。N就是patch的数目，决定了整个序列的长度。





Transformer部分的具体结构与NLP中的类似，主要经过多头自注意力（MSA）和MLP的处理，表示为：


MHSA是Transformer的核心，研究发现并行训练多个值的向量并将其拼接在一起得到最终输出比使用注意力函数得到输出效果更好，即使用多个Q-K-V矩阵并将这多个矩阵的输出拼接在一起，表示为：

由于每个头中向量的维度是等分的，因此使用多头注意力的计算代价与使用单头注意力的计算代价没有太大区别。

Section IV Experiments

Datasets
数据集为腹部CT数据集 来自MICCAI2015 你包含30例腹部CT扫描图像 
 输入图像大小为224 * 224。

Comparison Results

本文在多器官数据集上进行了测试，对比网络包括：VNet,DARR,UNet，AttnUNet,TransUNet 
Table 1展示了对比结果，可以看到本文的网络达到了最佳的性能，DSC和HD指标分别为78.09%和26.83%，其他网络结构的DSC指标均低于77.5%，HD指标均在31.55以上，证明了本文模型的有效性。



Fig 2展示了不同网络的分割细节，可以看到Attn-UNet存在着过分割的问题，而UNet则存在分割效果不佳的问题，主要是由于卷积的局限性和缺乏全局特征提取能力导致的，而本文对器官边界有更好的分割性能，尤其在一些小器官的分割上面，可以看到 本文的TransClaw UNet有效的学习了全局上下文信息，从而可以实现更精准的医学图像分割。




Ablation Experiments

本文还进行了消融实验来探究输入图像尺寸、skip connection数目、patch大小等对网络性能的影响。
Input size




当输入图像较大可以包含更多的像素信息，本文对比了输入尺寸为224和512时的模型性能，batch size分别为24和6，从table 2的结果可以看到，图像分辨率增加时DSC会显著增加。主要是因为随着图像分辨率的提升，patch size不变，序列长度增加，从而可以学习图像中更加详细的全局信息，并转换为更有效的编码结果，从而提升分割性能。但是图像分辨率的增加也需要更多的计算资源和更长的训练时间。
**number of skip connections:**

本文将跳跃连接的数目设置为0,1,2,3，fig 3展示不同跳跃连接数下的分割的DSC指标啊，可以看到连接数越少分割性能越差，因为skip connection可以有效结合编码、解码和上采样部分的信息，使得模型具有鲁棒性获得更好的性能。值得注意的是，skip connection == 0时Arota和Pancreas器官的DSC下降了10.8%和10.16%。


**patch size and sequence length：**

patch size会影响序列长度，而序列长度则会影响图像特征的提取。patch size较大，序列较短，能学习到的有效信息较少，影响了全局空间特征的获取；patch size较小，序列长度增大，但需要学习的参数迅速增长，极大的占用了计算资源。因此本文使用了16*16的patch size，此时DSC=78.09%,PATCH SIZE=24 DSC=77.27但是增长有限，内存也指数增长。

Section V Conclusion

本文提出的Trans Claw UNet有效结合了CNN和Transformer，借助多头自注意力有效的捕捉了全局上下文信息，弥补了卷积的局限性；在编码部分将二者结合取得了很好的分割效果，消融实验的结果表明本文的网络具有良好的泛化性能。