摘要:卷积神经网络(CNNs)的最新进展极大地促进了医学图像分割——这是众多临床需求的先决条件。然而,它在建模显式长程关系方面表现出普遍的局限性,而现有的解决方案,即通过构建深度编码器和积极的下采样操作,导致冗余的深化网络和局部细节的丢失。因此,分割任务等待更好的解决方案,以提高对全局上下文建模的效率,同时保持对低级细节的有力把握。在本文中,我们提出了一种新的并行分支架构TransFuse来应对这一挑战。TransFuse以并行的方式将Transformers和CNNs结合在一起,可以以更浅的方式有效地捕捉全局依赖性和低级空间细节。此外,还提出了一种新的融合技术——BiFusion模块,以有效地融合两个分支的多层次特征。大量实验表明,TransFuse在2D和3D医学图像集(包括息肉、皮肤病变、髋关节和前列腺分割)上都获得了最新的最先进的结果,显著降低了参数并提高了推理速度。
1 介绍
卷积神经网络(CNNs)在许多医学图像分割任务中获得了无与伦比的性能[9,12],如多器官分割、肝脏病变分割、大脑3D MRI等,因为它被证明通过端到端训练网络来构建分层任务特定特征表示。尽管基于CNN的方法取得了巨大成功,但其在获取全球背景信息方面缺乏效率仍然是一个挑战。感知全局信息的机会等于效率的风险,因为现有的工作是通过产生非常大的感受野来获得全局信息的,这需要连续的下采样和叠加足够深的卷积层,直到足够深。这带来了几个缺点:1)非常深的网络的训练受到特征重用递减问题的影响[23],其中低级别的特征被连续的乘法洗去;2) 随着空间分辨率逐渐降低,对密集预测任务(例如逐像素分割)至关重要的局部信息被丢弃;3) 用小的医学图像数据集训练参数重的深网往往是不稳定的,并且容易过拟合。一些研究[29]使用非局部自我注意机制来模拟全局上下文;然而,这些模块的计算复杂度通常相对于空间大小呈二次方增长,因此它们只能适当地应用于低分辨率的特征图。
Transformer最初用于对NLP任务中的序列到序列预测进行建模[26],最近在计算机视觉界引起了极大的兴趣。[7]中提出了第一个用于图像识别的纯基于自注意的视觉变换器(ViT),该变换器在ImageNet[6]上获得了有竞争力的结果,前提是在大型外部数据集上进行预训练。
SETR[32]在传统的基于编码器的网络中用转换器代替编码器,以成功地在自然图像分割任务中实现最先进的(SOTA)结果。虽然Transformer擅长对全局上下文进行建模,但它在捕捉细粒度细节方面存在局限性,尤其是对于医学图像。我们独立发现,由于在建模局部信息时缺乏空间归纳偏差,类似SETR的纯基于变换器的分割网络产生了不令人满意的性能(也在[4]中报道)。
为了享受两者的好处,已经努力将CNN与Transformer相结合,例如,TransUnet[4],它首先利用细胞神经网络提取低级特征,然后通过转换器对全局交互进行建模。TransUnet结合了跳跃连接,在CT多器官分割任务中创造了新的记录。然而,过去的工作主要集中在用变换器层代替卷积或以顺序的方式堆叠两者。为了进一步释放CNNs和transformer在医学图像分割中的力量,在本文中,我们提出了一种不同的架构——TransFuse,它并行运行基于浅层CNN的编码器和基于transformer的分割网络,然后是我们提出的BiFusion模块,将两个分支的特征融合在一起,共同进行预测。TransFuse具有以下优点:1)既可以有效地捕捉低级空间特征,又可以有效地捕获高级语义上下文;2) 它不需要很深的网,这缓解了梯度消失和特征减少的重用问题;3) 它大大提高了模型大小和推理速度的效率,不仅可以在云上部署,还可以在边缘部署。据我们所知,TransFuse是第一个合成CNN和Transformer的并行分支模型。实验证明,与其他竞争性SOTA作品相比,其性能优越。
2 提出的方法
如图所示,TransFuse由两个不同处理信息的平行分支组成:1)CNN分支,它逐渐增加感受野并从局部到全局对特征进行编码;并从局部到全局对特征进行编码;2) Transformer分支,从全局自关注开始,并在最后恢复局部细节。从两个分支提取的分辨率相同的特征被输入到我们提出的BiFusion模块中,其中应用自注意和双线性Hadamard乘积来选择性地融合信息。然后,将多级融合的特征图进行组合,以使用门控跳过连接生成分割[20]。那里
所提出的并行分支方法的两个主要优点是:首先,通过利用CNN和Transformers的优点,我们认为TransFuse可以在不构建非常深入的网络的情况下捕获全局信息,同时保持对低级别上下文的敏感性;其次,我们提出的BiFusion模块可以在特征提取过程中同时利用细胞神经网络和变换器的不同特征,从而使融合表示强大而紧凑。
Transformer分支
Transformer分支的设计遵循典型的编码器-解码器架构。具体地,首先将输入图像x∈R^(H×W×3)均匀地划分为N=H/S×W/S块,其中S通常设置为16。然后,将补丁展平并传递到输出维度为D0的线性嵌入层中,获得原始嵌入序列e∈RN×D0。为了利用空间先验,将相同维度的可学习位置嵌入添加到e中。得到的嵌入z0∈R^N×D0是Transformer编码器的输入,该编码器包含L层多头自注意(MSA)和多层感知器(MLP)。我们强调,自注意(SA)机制是Transformer的核心原理,它通过在每一层全局聚合信息来更新每个嵌入补丁的状态:
Wqkv∈RD0×3Dh是投影矩阵,向量zi∈R1×D0,qi∈R1 X Dh分别是z和q的第i行。
MSA是SA的扩展,它将多个SA连接起来,并将潜在维度投影回RD0,MLP是密集层的堆栈(有关MSA和MLP的详细信息,请参阅[7])。将层归一化应用于最后一个变换器层的输出,以获得编码序列zL∈R^N×D0。对于解码器部分,我们使用渐进上采样(PUP)方法,如SETR[32]中所述。具体来说,我们首先将zL重塑为t0∈R H 16×W 16×D0,它可以被视为具有D0通道的2D特征图。然后,我们使用两个连续的标准上采样卷积层来恢复空间分辨率,其中我们分别获得t1∈R H 8×W 8×D1和t2∈R H.4×W 4×D2。保存不同尺度t0、t1和t2的特征图,以便与CNN分支的相应特征图进行后期融合。
CNN分支
传统上,特征被逐步下采样到H32×W32,并且在深层细胞神经网络中使用数百层来获得特征的全局上下文,这导致非常深层的模型耗尽资源。
考虑到Transformers带来的好处,我们从原始CNN管道中删除了最后一个块,并利用Transformer分支来获取全局上下文信息。这不仅为我们提供了一个更浅的模型,而且还保留了更丰富的本地信息。例如,基于ResNet的模型通常有五个块,每个块对特征图进行二倍的下采样。我们将第4个(g0∈R^H /16×W /16×C0)、第3个(g1∈R^H/8×W/8×C1)和第2个(g2∈R^H/4×W /4×C2)块的输出与Transformer的结果进行融合(图1)。此外,我们的CNN分支是灵活的,任何现成的卷积网络都可以应用。
BiFusion模块
为了有效地结合来自CNNs和Transformers的编码特征,我们提出了一个新的BiFusion模块(见图1),该模块结合了自注意和多模态融合机制。包含的操作有:
通道注意力、空间注意力、3x3卷积、残差连接。具体地,我们通过以下操作获得融合特征表示f^i,i=0,1,2:
是Hadamard乘积,Conv是3x3卷积层。
通道注意力被实现为[10]中提出的SE块,以促进来自Transformer分支的全局信息。
来自CBAM[30]块的空间注意力被用作空间滤波器,以增强局部细节并抑制不相关区域,因为低级别的CNN特征可能是噪声的。然后,Hadamard产品对来自两个分支的特征之间的细粒度交互进行建模。
最后,将交互特征b^i和参与特征t^i、g^i连接并通过残差块。所得到的特征f^i有效地捕获了当前空间分辨率的全局和局部上下文。为了生成最终分割,使用注意力门控(AG)跳跃连接[20]对f^i进行组合,其中我们有,如图1所示。
损失函数
使用加权IoU损失和二进制交叉熵损失对整个网络进行端到端训练
其中边界像素接受更大的权重分割预测是由一个简单的头生成的,它直接将输入特征图的大小调整为原始分辨率,并应用卷积层生成M个图,其中M是类的数量。根据[8],我们使用深度监督,通过额外监督变压器支路和第一个融合支路来改善梯度流。最终训练损失由其中α,γ,β是可调谐的超参数,G是基本事实。
3 实验和结果
数据采集
为了更好地评估TransFuse的有效性,考虑了四个具有不同成像模式、疾病类型、目标对象、目标大小等的分割任务:1)息肉分割,其中使用了五个公共息肉数据集:Kvasir[14]、CVC ClinicDB[2]、CVC ColonDB[24]、EndoScene[27]和ETIS[21]。采用与[8,11]中所述相同的分割和训练设置,即1450张训练图像仅从Kvasir和CVC ClinicDB中选择,798张测试图像来自所有五个数据集。在处理之前,每个图像的分辨率被调整为352×352,如[8,11]。2) 皮肤损伤分割,其中使用公开的2017年国际皮肤成像协作皮肤损伤分割数据集(ISIC2017)[5]。
ISIC2017提供2000张图像用于训练,150张图像用于验证,600张图像用于测试。根据[1]中的设置,我们将所有图像的大小调整为192×256。
3) 髋关节分割,共从医院收集641例病例,平均大小为2942×2449,像素间距为0.143mm4。每张图像由一名临床专家进行注释,并由两名专家进行双盲审查。我们将所有图像的大小调整为352×352,并以7:1:2的比例随机分割图像进行训练、验证和测试。4) 前列腺分割,其中使用来自医学分割十项[22]的体积前列腺多模态MRI。该数据集包含32名患者的多模态核磁共振成像,中位体积形状为20×320×319。根据[12]中的设置,我们将所有MRI切片重塑为320×320,并使用z评分归一化对每个体积进行独立归一化。
实施细节
TransFuse构建在PyTorch框架[16]中,并使用单个NVIDIA-100 GPU进行训练。根据经验,α、β和γ的值分别设置为0.5、0.3和0.2。采用了学习率为1e-4的Adam优化器并且所有模型都被训练了30个时期以及16个批量。
在息肉分割实验中,除了多尺度训练外,没有使用数据增强,如[8,11]所示。对于皮肤损伤和髋关节分割,在训练过程中应用了数据增强,包括随机旋转、水平翻转、颜色抖动等。发现7e-5的较小学习率对于皮肤损伤分割是有用的。最后,我们遵循nnU-Net框架[12],使用相同的数据增强和后处理方案来训练和评估我们的前列腺分割模型。由于所选的预训练数据集和分支主干可能会对性能产生不同的影响,因此提供了TransFuse的三种变体:1)更好地证明了我们方法的有效性和灵活性;2) 与其他方法进行公平比较。TransFuse-S分别以ResNet-34(R34)和8层DeiT-Smal(DeiT-S)[25]作为CNN分支和Transformer分支的主干来实现。类似地,TransFuse-L是基于Res2Net-50和10层DeiT-Base(DeiT-B)构建的,而TransFuse-L*使用ResNetV2-50和ViT-B[7]。请注意,ViTs和DeiTs具有相同的主干架构,它们主要在预训练的策略和数据集上有所不同:前者在ImageNet21k上训练,而后者在ImageNet1k上训练,数据扩充更重。
评估结果
TransFuse在2D和3D数据集上进行了评估,以证明其有效性。由于不同的医学图像分割任务用于不同的诊断或手术目的,我们遵循每个分割任务常用的评估指标来定量分析结果。TransFuse-S选定的可视化结果如图2所示。
图2 所有三项任务的结果可视化(最好用颜色查看)。每一行遵循地面实况(GT)和预测(Pred)的重复序列。
息肉分割结果
我们首先评估了我们提出的方法在息肉分割方面相对于各种SOTA方法的性能,即平均Dice(mDice)和平均联合交集(mIoU)。
如表1所示,我们的TransFuse-S/L在很大程度上优于基于CNN的SOTA方法。具体而言,TransFuse-S在未发现的数据集(ColonDB、EndoSene和ETIS)上实现了5.2%的平均mDice改进。与其他基于变压器的方法相比,TransFuse-L*在Kvasir和ClinicDB上也表现出卓越的学习能力,与TransUnet相比,mIoU增加了1.3%。此外,在配备Xeon(R)Gold 5218 CPU的RTX2080Ti上评估了参数数量和推理速度方面的效率。与现有的基于CNN的技术相比,TransFuse-S在仅使用26.3M个参数的情况下实现了最佳性能,相对于HarDNet MSEG(33.3M)和PraNet(32.5M)降低了约20%。
此外,由于我们提出的并行分支设计,TransFuse-S能够以98.7 FPS的速度运行,比HarDNet MSEG(85.3 FPS)和PraNet(63.4 FPS)快得多。类似地,与其他基于Transformerbased的方法相比,TransFuse-L*不仅获得了最好的结果,而且运行速度为45.3 FPS,比TransUnet快约12%。
皮肤损伤分割结果
ISBI 2017挑战赛根据ISIC 2017测试集上的Jaccard索引[5]对方法进行排名。在这里,我们使用Jaccard指数、骰子得分和像素精度作为评估指标。这个与主流方法的比较结果如表2所示。TransFuseS在Jaccard评分方面比之前的SOTA SLSDeep[19]好1.7%,没有任何预处理或后处理,并且在不到1/3的时期内收敛。此外,我们的结果优于采用预训练的R34作为主干的Unet++[33],并且具有与TransFuse-S相当数量的参数(26.1M与26.3M)。再次,结果证明了我们提出的架构的优越性。
表1:与以前的SOTA相比,息肉分割数据集的定量结果。[4]的结果是通过运行发布的代码获得的,我们实现了SETR-PUP意味着结果不可用。
与主流方法的比较结果如表2所示。TransFuseS在Jaccard评分方面比之前的SOTA SLSDeep[19]好1.7%,没有任何预处理或后处理,并且在不到1/3的时期内收敛。此外,我们的结果优于采用预训练的R34作为主干的Unet++[33],并且具有与TransFuse-S相当数量的参数(26.1M与26.3M)。再次,结果证明了我们提出的架构的优越性。
髋关节分割结果
表3显示了我们对髋关节分割任务的结果,该任务涉及三个人体部位:骨盆、左股骨(L-股骨)和右股骨(R-股骨)。由于轮廓在dianosis和THA术前计划中更为重要,我们使用Hausdorff距离(HD)和平均表面距离(ASD)来评估预测质量。与两种先进的分割方法[33,28]相比,TransFuse-S在两种指标上都表现最好,并显著降低了HD(与HRNetV2相比平均降低30%,与Unet++相比平均降低34%),这表明我们提出的方法能够捕捉更精细的结构并生成更精确的轮廓。
前列腺分割结果
我们将TransFuse-S与nnUNet[12]进行了比较,后者在前列腺分割挑战中排名第一[22]。我们遵循与公众相同的预处理、培训和评估方案可用的nnU-Net框架5,并在表4中报告5倍的交叉验证结果。我们可以发现,就平均骰子得分而言,TransFuse-S以很大的优势(+4.2%)超过了nnUNet-2d。与nnUNet-3d相比,TransFuse-S不仅实现了更好的性能,还将参数数量减少了~41%,吞吐量增加了~50%(在GTX1080上)。
消融研究
通过不同主干、组成和融合方案的不同设计选择,进行了消融研究,以评估并行分支内设计以及BiFusion模块的有效性。使用息肉的可见(Kvasir)和未可见(ColonDB)数据集,结果记录在平均骰子中。在表5中,通过将E.3与E.1和E.2进行比较,我们可以看到,将CNN和Transformer相结合会带来更好的性能。此外,通过比较E.3和E.5E.6,我们观察到并行模型的性能优于顺序模型。此外,我们使用与我们提出的E.6相同的并行结构和融合设置来评估双分支CNN模型(E.4)的性能。我们观察到,E.6在Kvasir和ColonDB中分别比E.4高2.2%和18.7%,这表明CNN分支和transformer分支是互补的,从而获得更好的融合结果。最后,对另一个包含级联和残差块的融合模块与我们提出的BiFusion模块(E.5和E.6)进行了性能比较。在相同的主干和组成设置下,E.6与BiFusion实现了更好的结果。在ISIC2017上进行的额外实验如表所示。6验证了BiFusion模块的设计选择,我们发现每个组件都显示出其独特的优势。
4 结论
在本文中,我们提出了一种新的策略,将Transformers和CNNs与后期融合相结合,用于医学图像分割。由此产生的架构TransFuse利用了CNNs在建模空间相关性方面的归纳偏差,以及Transformers在建模全球关系方面的强大能力。
TransFuse在各种分割任务上实现了SOTA性能,同时在参数和推理速度上都非常高效。我们希望这项工作能够为使用基于转换器的架构带来一个新的视角。
未来,我们计划提高香草变压器层的效率,并在其他医学相关任务上测试TransFuse,如里程碑式检测和疾病分类。