3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation

加粗样式

目录

  • 翻译
    • 摘要
    • 1.引言
    • 2.所提出的方法
        • 图一
    • 3.实验与结果
        • 图二
        • 表一
        • 表二
        • 表三
        • 表四
        • 表五
        • 表六
    • 4.结论
  • 笔记

原文地址:
https://paperswithcode.com/paper/transfuse-fusing-transformers-and-cnns-for

翻译

摘要

医学图像分割-许多的临床需求必备条件-卷积神经网路的最新进展使其得到了显著的发展。但是,它在建模显示远程关系(modeling explicit long-range relation)存在普遍的限制,现有的解决方法采取深度编码器和下采样操作(deep encoders along with aggressive downsampling operations),导致冗余的深层次网络和局部细节的损失。因此,分割任务正在等待一个更好的解决办法,来提高全局上下文建模的效率,同时保持对低层细节(low-level details)的强大掌握。在本文中,我们提出了一个新的分支并行(parallel-in-branch)结构,TransFuse,来解决这个问题。TransFuse以并行的方式结合了Transformers和CNNs,其中全局依赖和低层空间信息都可以以浅(shallower)的方式被有效获取。此外,还提出了一种新的融合技术,BiFusion模块,可以有效的融合来自不同分支的多层次特征(multi-level features)。大量实验表明,TransFuse在息肉、皮肤病变、髋部和前列腺分割在内的2D、3D医学图像集上都取得了最新的研究结果,显著降低了参数并提高了推理速度
关键词:医学图像分割;Transformers;卷积神经网络;Fusion

1.引言

卷积神经网络(CNNs) 在许多医学图像分割任务[9,12]上能够得到非并行性能,例如多器官分割,肝脏病变分割,3D大脑MRI等,结果表明它有强大的能力通过端对端的训练网络建立分层的特定任务特征表示。尽管基于CNN方法学的巨大成功,它在获取全局文本信息的性能缺失依然是一个问题。感受全局信息的几率与效率的风险相等,因为现存的工作通过产生很大的视野域(receptive field)获取全局信息,其中视野域需要不断的下采样和叠加卷积层直到深度足够。这带来了一些缺点:1)训练很深的网络被递减特征重用问题(the diminishing feature reuse problem)[23]影响,其中低层特征被连续乘法(consecutive multiplication)冲刷掉。2)随着空间分辨率的减小,对例如像素级(pixel-wise)分割的密集预测任务很重要的局部信息,被丢弃。3)带有小医学图像数据的训练参数深层网( training parameter heavy deep nets )更容易不稳定和过拟合。一些研究[29]使用非局部自我注意力机制建模全局文本;但是,这些方法的计算复杂度通常与空间大小成二次增长,因此它们只适用于低分辨率的maps。
Transformer,最初在NLP任务[26]中用于建模序列对序列预测,最近在计算机视觉领域引起巨大兴趣。基于视觉变换器(ViT)的第一个纯自注意力机制在[7]中提出,用于图像识别,在较大的外部数据库上进行预训练的前提下,在ImageNet[6]的结果具有竞争性。在传统的基于编码器-解码器的网络上,SETR[32]把编码器替换为变换器,成功达到了自然图像分割任务上最好的结果。虽然Transformer善于建模全局文本,在获取细粒度(fine-grained)细节上有局限性,特别是用于医学图像。我们独立的发现,由于在建模局部信息(在[4]中报道)上空间诱导偏差(inductive-bias)的缺失,基于分割网络的类似于SETR的纯变换器的分割网络的性能不是很满意。
为了同时利用两者的优点,人们努力把CNNs和Transformer结合,例如TransUnet[4],它首先利用CNNs提取低层特征然后通过变换器建模全局交互。通过结合跳跃连接,TransUnet在CT多器官分割任务中达到了新纪录。但是,过去的工作主要集中在把卷积替换为变换层或者以顺序方式堆叠两个变换层。为了更进一步利用CNN和变换器在医学图像分割的能力,本文中,我们提出了一个不同的结构-TransFuse,它并行运行基于CNN的编码器和基于变换器的分割网络,然后是我们提出的BiFusion模块,其中来自两个分支(branch)的特征被融合在一起,共同进行预测。TransFuse具有以下一些优点:1)可以有效获取底层空间特征和高层语义文本;2)不需要很深的网,缓解了梯度消失(gradient vanishing)和特征递减重用问题;3)很大程度上提高了模型大小和推理速度的效率,能够云端和边缘(edge)部署。据我们所知,TransFuse是第一个综合CNN和Transformer的分支内并行的模型。实验表明,该性能优于其他竞争的最好的工作。

2.所提出的方法

如图1所示,TransFusion包括两个并行的分支,处理不同的信息:1)CNN分支,逐渐增加视野域和从局部到全局对特征编码;2)Transformer分支,从全局自注意开始,并在最后恢复局部细节。从两个分支提取的相同分辨率的特征被输入到我们所提出的BiFusion模块,其中自注意和双线Hadamard乘积被应用于选择性地融合信息。之后,多层被融合的特征映射图结合起来,使用门控跳过连接[20]生成分割。对于所提出的并行分支方法有两个主要的优点:首先,通过利用CNNs和Transformer的优点,我们认为TransFuse能够在没有创建很深的网络时可以获取全局信息,同时保持对低层次文本的敏感度;其次,我们所提出的BiFusion模块在特征提取时,能够同时利用CNNs和Transformer的不同特性,因此使得融合表示更加强大和紧凑。

图一

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第1张图片
图一:TransFuse的概述(最佳彩色试图):两个并行的分支-CNN(右下)和transformer(左)由我们提出的BiFusion模块融合。

  Transformer分支:Transformer分支的设计遵循典型的编码器-解码器结构。具体来说,输入图像x∈RH×W×3首先被平均分为N=(H/S)×(W/S)patch,其中S通常被设置为16。patch被打平后送入输出维度为D0的线性嵌入层,得到原始的嵌入序列e∈RN×D0。为了利用空间先验(the spatial prior),在e上添加一个相同维度的可学习的位置嵌入。结果嵌入z0∈RN×D0是变换器编码器的输入,其中包含L层多头自注意(MSA)和多层感知器(MLP)。我们强调了自注意(SA)机制,这是变换器的核心原理,它通过增加每一层中全局聚合信息(aggregating information globally),更新了每一个嵌入patch的状态:
3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第2张图片
其中[q,k,v]=zWqkvWqkv∈RD0×3Dh是投影矩阵(projection matrix),向量zi∈R1×D0qi∈R1×Dh分别是z和q的第i行。MSA是SA的扩展,它连接多层SA并且将潜在维度投影回RD0,MLP是密集层的堆栈(有关MSA和MLP的详细信息请参阅[7])。对最后一个变换层的输出进行层归一化(Layer normalization)处理,得到编码后的序列zL∈RN×D0。对于解码部分,我们使用渐进上采样(PUP)(progressive upsampling)方法,如SETR[32]。具体的来说,我们首先把zL塑形为t0∈RH/16×W/16×D0,可以看作一个具有D0通道的2D特征图。然后我们使用两个连续的标准上采样卷积层,用于恢复空间分辨率,其中我们分别得到t1∈RH/8×W/8×D1t2∈RH/4×W/4×D2。保存t0t1t2不同规模的特征图,与CNN分支相应的特征图进行后期融合。

  CNN分支:传统上,特征被逐步下采样到H/32×W/32,在深度CNN中使用数百层来获取全局上下文特征,这使得很深的模型消耗资源。考虑到变换器带来的好处,我们从最初CNNpipeline中移除了最后一块,相反,利用变换分支的优点得到全局上下文信息。这不仅使我们得到一个更浅的模型,而且保持了更丰富的局部信息。例如,基于ResNet的模型通常有5个块,每一个块特征图都进行2倍的下采样。我们把来自于第4层(g0∈RH/16×W/16×C0),第3层(g1∈RH/8×W/8×C1)和第2层(g2∈RH/4×W/4×C2)块的输出与来自变换器(图1)的结果融合。而且,我们的CNN分支更加灵活,可以应用于任何现成的卷积网络。

  BiFusion模块:为了有效结合CNN和变换器编码后的特征,我们提出了一个新的BiFusion模块(如图1),它包含自注意和多模式融合机制(multi-modal fusion mechanisms)。具体地说,通过接下来的操作,我们得到了融合特征表示fi,i=0,1,2:
在这里插入图片描述
其中Wi1∈RDi×Li,Wi2∈RCi×Li,|⊙|是Hadamard乘积,Conv是一个3×3的卷积层。channel attention按照[10]中提出的SE-Block实现,用于促进来自变换器分支的全局信息。因为低层的CNN特征可能是噪音,采用CBAM[30]块的spatial attention作为空间滤波器(spatial filters),用于增强局部细节和抑制不相关的区域。然后Hadamard乘积对来自两个不同分支的特征间细粒度交互进行建模。最后,将交互特征b^i 与参与特征(attended features)t^ig^i 连接,然后通过一个Residual块。生成的特征fi有效获取了当前空间分辨率的全局和局部文本。为了产生最终的分割,所有的fi通过使用注意门(AG)跳跃连接[20]进行组合,其中我们有f^i+1=Conv([Up(f^i),AG(fi+1,UP(f^i))]),f^0=f0,如图1。

  损失函数:采用加权IoU损失和二元交叉熵损失L=LwIoU+Lwbce对全网络进行端到端的训练,其中边界像素接收较大的权重[17]。分割预测由一个simple head产生,可以直接把输入特征图调整为原始分辨率,并且应用卷积层产生M个映射maps,其中M是类的数量。在[8]之后,我们通过额外监督变换器分支和第一个融合分支,使用深度监督(deep supervision)来提高梯度流(gradient flow)。最后的训练损失等于=αL(G,head(f^2))+γL(G,head(t2))+ βL(G,head(f0)),其中α,,γ, β是可调整的超参数,G是groundtruth。

3.实验与结果

数据采集:为了更好评估TransFuse的有效性,我们考虑了4种带有不同成像模式,疾病类型,target object,目标尺寸等的分割任务:1)息肉分割,其中使用了五个公开的息肉数据集:Kvasir[14],CVC-ClinicDB[2],CVC-ColonDB[24],EndoScene[27]和ETIS[21]。采用了[8]种描述的相同的分割和训练设置,即从Kvasir和CVC-ClinicDB中单独选择1450张训练图像,同时从所有5个数据集中选择798张测试图像。处理前,每个图像的分辨率调整为352×352,如[8,11]。2)皮肤病变检测,其中使用了公开可用的 2017 International Skin Imaging Collaboration皮肤病变分割数据集(ISIC2017[5])。ISIC2017提供了2000张图像用于训练,150张用于验证,600张用于测试。按照[1]中的设置,我们把所有图像调整为192×256。3)髋部分割,其中共从医院收集了641例病例,平均大小为2942×2449,像素间距为0.143mm4。每张图像都由一位临床专家进行注释,并由两位专家进行双盲检查。我们把所有图像调整为352×352,并且按7:1:2的比例随机分割图片用于训练,验证和测试。4)前列腺分割,其中使用了来自 Medical Segmentation Decathlon[22]的volumetric Prostate Multi-modality MRIs。数据库包含来自32个病人的多模式MRIs,平均体积形状为20×320×319。按照[12]中的设置,我们把所有的MRI切片调整为320×320,使用z-score标准流对每一个体积(volume)进行独立归一化。

运行细节:TransFuse创建在PyTorch框架[16]中,使用单个NVIDIA-A100 GPU进行训练。α,,β和γ的值根据经验设置为0.5,0.3,0.2。采用学习率为1e-4的Adam优化器,除非另有规定,否则所有模型都是训练30个epoch,batch size为16。
在息肉分割实验中,除了多尺度训练(multi-scale training),没有使用数据增强,如[8,11]。对于皮肤病变和髋部分割,在训练期间应用了包含随即旋转,水平翻转,颜色抖动(color jittering)等数据增强。对于皮肤病变分割发现较小的学习率7e-5更有用。最后我们遵循nnU-Net框架[12],使用相同的数据增强和后处理方案(post-processing scheme),对我们的前列腺分割模型进行训练和评估。由于选定的预训练数据集和分支backbones可能对性能产生不同的影响,因此提供了TransFuse的三个变体:1)更好的展示了我们方法的有效性和灵活性;2)与其他方法进行了公平的比较。TransFuse-S使用ResNet-34 (R34)和8-layer DeiT-Small (DeiT-S)[25]分别作为CNN分支和变换器分支的backbone。同样的,TransFuse-L基于Res2Net-50和10-layer DeiT-Base (DeiT-B)建立,同时TransFuse-L*使用ResNetV2-50和ViT-B[7]。值得注意的是,ViTs和DeiTs拥有相同的backbone结构,他们主要在预训练策略和数据集上不同:前者在ImageNet21k上训练,而后者在ImageNet1k上训练并进行更大的数据增强。
评估结果:TranFuse同时在2D和3D数据库上评估,来证明有效率。由于不同的医学图像分割任务用于不同的诊断或者手术目的,我们遵循每一个分割任务常用的评估指标来定量分析结果。选定的TransFuse-S可视化结果如图2所示。

图二

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第3张图片
图二:在三个任务(最好由彩色查看)上的可视化结果。每一行都遵循ground truth(GT)和预测(Pred)的重复序列。

  息肉分割的结果:在平均Dice(mDice)和平均Intersection-Over-Union(mIoU)的角度,首先我们在息肉分割上与一些系列最新的方法一起评估了我们提出的方法的性能。如表一,我们的TransFuse-S/L远远超过基于CNN的最新的方法。具体地说,TransFuse-S在未发现的数据库(ColonDB, EndoSene and ETIS)上实现了平均mDice5.2%的改进。与其他基于变换器的方法比较,TransFuse-L*在Kvasir和ClinicDB上也展现了优越的学习能力,观察到与TransUnet比较,在mIoU上提升了1.3%。此外,就参数数量和推理速度的效率而言,是在带有Xeon® Gold 5218 CPU的RTX2080Ti上评估的。与之前基于CNN的技术比较,TransFuse-S在仅使用26.3M参数的情况下达到了最好的性能,与HarDNet-MSEG(33.3M)和PraNet (32.5M)相比减少了大约20%。而且,多亏我们提出的并行分支设计,TransFuse-S能够以98.7FPS的速度运行,比HarDNet-MSEG(85.3FPS)和PraNet (63.4FPS)更快。同样的是,TransFuse-L*与其他基于变换器的方法相比,不仅能够达到最好的结果而且运行速度达到45.3FPS,比TransUnet快了12%左右。

表一

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第4张图片
表一:在息肉分割数据库上与其他之前最新技术比较的定量结果。[4]的结果通过运行发布的代码得到,我们实现了SETR-PUP。‘_’表示结果不可用。

  皮肤病变分割的结果:ISBI2017挑战根据ISIC2017测试集上的Jaccard Index[5]对方法进行排名。这里,我们使用Jaccard Index,Dice分数和像素级准确度作为评估指标。与最新的方法的比较结果如表二所示。在没有预处理或者后处理并且在少于1/3epoch收敛的情况下,TransFuse-S在Jaccard分数上比之前最新技术SLSDeep[19]更好大约1.7%。并且,我们的结果超过Unet++[33],后者采用预训练的R34作为backbone且与TransFuse-S(26.1Mvs26.3M)的参数数量相当。我们的结果再次证明了我们提出的结构的优越性。

表二

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第5张图片

表二:在ISIC2017测试集上的定量结果。backbone的结果使用了在ImageNet上预训练的权重。

  髋部分割的结果表三展示了在髋部分割任务的结果,涉及了三个人体部分:骨盆、左股骨(L-股骨)和右股骨(R-股骨)。由于轮廓THA术前计划中更重要,我们使用Hausdorff Distance(HD)和平均表面距离(ASD)来评估预测质量。与两个先进的分割方法[33,28]对比,TransFuse-S在两个指标上都达到了最好并且明显减少了HD(与HRNetV2比较平均降低30%,与Unet++相比平均降低了34%),表明我们所提出的方法能够获取更精细的结构和产生更加精确的轮廓。

表三

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第6张图片
表三:在内部髋部数据集上的结果。所有的模型都使用了在ImageNet上预训练的backbone,并且具有相似的大小(约26M)。HD和ASD以mm毫米为单位。

  前列腺分割的结果:我们比较了TransFuse-S和nnUNet[12],后者在前列腺分割挑战[22]中排名第一。我们遵循公共可用的nnUNet框架( https://github.com/MIC-DKFZ/nnUNet)的相同的预处理,训练和评估方案,并在表四中报告了5-fold交叉验证结果。我们发现TransFuse-S在平均dice得分方面远远超过了nnUNet-2d(+4.2%)。与nnUNet-3d相比,TransFuse-S不仅达到了更好的性能,也减少了参数数量约41%,并且增加了吞吐量约50%(在 GTX1080上)。

表四

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第7张图片
表四:在前列腺MRI分割的定量结果。PZ,TZ代表两个标签类(外围和过渡区),性能(PZ,TZ和均值)由dice分数衡量。

消融实验:通过改变不同backbone、成分和融合方案的设计选择,实施消融实验,用来评估并行分支设计和BiFusion模块的有效性。使用了来自息肉的可视(Kvasir)和不可视的(ColonDB)数据库,结果以平均Dice记录。在表五中,通过将E.3与E.1和E.2比较,我们可以看到结合CNN和变换器能达到更好的性能。此外,通过将E.3与E.5和E.6比较,我们可以观察到并行的模型比顺序模型的性能要好。而且,通过使用我们提出的E.6相同的并行结构和融合设置,我们评估了双分支CNN模型(E.4)的性能。我们观察到E.6在Kvasir上优于E.4达2.2%,在 ColonDB上优于18.7%,表明CNN分支和变换器分支是互补,可得到更好的融合结果。最后,将另一个融合模块(包括拼接和residual块)与我们提出的BiFusion模块(E.5和E.6)进行性能比较。在相同backbone和成分设置下,E.6结合BiFusion的效果更好。表六给出了在ISIC2017上进行的其他实验,以验证BiFusion模块的设计选择(design choice),从中我们发现每一个成分都显示出独特的优点。

表五

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第8张图片
表五:在并行分支设计上的消融实验。Res:Residual。

表六

3.翻译和笔记TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation_第9张图片
表六:在BiFusion模块上的消融实验。Res:Residual;TFM:Transformer;Attn:Attention。

4.结论

在本文中,我们提出了一个新的策略,结合Transformers和CNNs与后期融合(with late fusion)用来医学图像分割。由此产生的框架,TransFuse,利用了CNN在建模空间相关性的归纳偏置(inductive bias)和Transformers在建模全局关系上的强大能力。TransFuse在各种分割任务上达到了最新的技术,同时在参数和推理速度上都很有效。我们希望这项工作能够为基于transformer的结构带来新的视角。在未来,我们计划改善vanilla transformer层的效率,并在其他医学相关任务测试TransFuse,比如landmark检测和疾病分类。

Acknowledgement:我们非常感谢Weijun Wang, MD,Zhefeng Chen,MD, Chuan He,,MD,Zhengyu Xu, Huaikun Xu担任我们髋关节分割项目的医学顾问。

笔记

你可能感兴趣的:(论文学习,计算机视觉,深度学习,人工智能)