基于Transformer的核磁共振肿瘤图像分割研究

吕嫄

(芜湖职业技术学院 基础教学部,安徽 芜湖,241000)

核磁共振成像可以很好地显示人体的组织信息,用于区分肿瘤是否为癌性以及癌细胞是否转移。因此,核磁共振成像可以辅助医务人员规划肿瘤治疗方案。近年来,随着深度学习技术在图像处理领域的迅速发展,运用深度学习技术对肿瘤病人的核磁共振影像进行自动分割可以更加快速准确地为医生提供诊断依据。

基于深度学习技术的卷积神经网络被广泛应用于医学图像分割领域[1-3]。U-net是最早提出U形分割架构的卷积神经网络,其设计思想得到了广泛的应用[4-6]。在卷积神经网络中,有很多工作都证明了多尺度特征对于图像分割任务非常重要[7-8]。池化操作会损害这种多尺度的信息,空洞卷积解决了这个问题[9]。但是这种卷积会造成信息的不连续,因此,本文提出一种不使用卷积的模型。

Transformer最早是用于解决自然语言处理领域中无法并行训练的问题。直到视觉Transformer巧妙地将其应用于图像处理,弥补了卷积神经网络的一些缺点,推进了医学图像处理领域的发展[10-11]。Swin-Transformer在视觉Transformer基础上提出了滑动窗口算法,降低了计算复杂度[12]。nnFormer和VT-UNet是两个完全基于Transformer架构的3D分割网络[13-14]。TransUNet是第一个提出结合卷积和Transformer的医学图像分割网络,但是它只用了一个卷积层来提取Transformer的特征[15]。ViTAE通过使用不同膨胀率的多重卷积向下采样,将输入图像嵌入到具有多尺度上下文的令牌中[16]。UNETR,TransBTS都是基于视觉Transformer架构的3D医学图像分割网络,计算复杂度高[17-18]。Cross-view Transformer通过在全局池化层之前链接特征图,在视图之间传输信息,在胸部X射线数据集上取得优越的性能[19]。FCT(fully convolutional Transformer)是第一个用于医学影像应用的全卷积Transformer模型。它利用Transformer提取图像的长距离语义依赖,然后利用卷积提取图像的局部属性,在ISIC 2017数据集上取得较好性能[20]。D-former 提出了一种膨胀Transformer,在不增加计算成本的情况下扩大感受野,提高了模型对于3D医学图像的分割性能。DS-TransUNet 将分层Swin Transformer 融合到U形结构的编码器与解码器中,有效地提取了非局部依赖和多尺度上下文信息,提高了不同医学图像的语义分割质量。GFENet以视觉Transformer作为基础单元构建了一个金字塔结构的特征提取器,通过线性预测和局域解码对提取的特征进行局部增强,并采用全局注意力将位置信息嵌入特征,在CVC-ColonDB数据集上获得了优于其他高性能分割模型的结果。经过综合分析各模型的优缺点,本文提出了一种完全基于Transformer的分割网络,它以Transformer的变体Swin Transformer作为基础模块,不仅可以准确地分割三维核磁共振的脑肿瘤图像,而且参数量少,计算复杂度低。

1 方法

1.1 数据集

本文为了验证所提出网络结构的有效性,使用了公共的脑肿瘤分割数据集BraTS2021进行了实验验证[21]。该数据集一共包含1 251例患者的三维核磁共振影像。本文对所有样本进行随机抽样,将834例样本作为训练集,208例样本作为验证集,209例样本作为测试集。其中,每例样本由4种模态(Flair,T1,T1CE,T2)的3D图像以及1个共享标签组成。

每个模态图像都是磁共振成像,长×宽×高为240×240×155,单个的体积像素是1 mm3。如图1所示,T1(T1-weighted imaging)通过对人体进行磁共振成像,主要用于显示解剖结构;T1CE(T1-weighted imaging with contrast enhancement)在T1序列的基础上加入了造影剂,主要用于显示肿瘤和炎症等异常组织;T2(T2-weighted imaging)使用对T2信号敏感的磁共振成像方法,用于显示病变的水肿或液体;Flair(Fluid Attenuated Inversion Recovery)通过将T2序列与T1序列结合,用于显示肿瘤、炎症和其他异常组织。

图1 数据集4种模态示例Fig.1 Four modal examples for datasets

1.2 网络结构

本文提出的分割网络的整体结构如图2所示,它主要由收缩模块、扩张模块和跳转连接组成。其中收缩模块用于提取和融合输入图像的特征;扩张模块用于恢复特征图的原始分辨率;跳转连接有助于恢复图像的细粒度细节。当输入一张图像后,经过多个编码层得到不同尺度的立体特征图(图2中输入分支中的立方体),并将编码后的特征图直接输入扩张模块进行解码得到新的特征图(图2中输出分支中的立方体)。同时,收缩模块通过跳转注意力与跳转连接向扩张模块传递多尺度信息。

图2 分割网络整体结构Fig.2 The overall structure of segmented network

1.2.1 收缩模块

收缩模块由3个独立的Swin Transformer结构组成,对应图2中的编码层。每一个Swin Transformer组成相同并且与传统的结构相同,如图3所示,W-MSA表示窗口注意力,SW-MSA表示滑动窗口注意力,MLP表示多层感知机。Swin Transformer的第一层是归一化层(LN),第二层是多头自注意力层,第三层是第二个归一化层,第四层是多层感知机,对每个像素点分类。该模块的计算步骤如下:

图3 单个Swin Transformer模块结构Fig.3 The architecture of a Swin Transformer

(1)

(2)

(3)

(4)

式中:l为第l层;z为图4中不同模块的输出。每一层的注意力计算公式具体如下:

图4 训练过程Fig.4 The training process

Q=ΩqX+βqlT
K=ΩkX+βklT
V=ΩvX+βvlT

(5)

式中:Q、K和V分别为查询向量、键向量和值向量;βq、βk、βv∈RD,βq、βk和βv为可选择偏置项;Ωq∈Rd×dq,Ωk∈Rd×dk,Ωv∈Rd×dv;lT为偏置项的转置。

图像数据通过归一化层后削弱噪声影响并且减小数据间的差距,这使得传入多头自注意力层的数据能够在分头后每一部分都具有源数据的特征,进而保证多层感知机输出的特征是合理的,并且模型整体训练速度不会降低。每一个Swin Transformer块输出特征的尺寸是递减的,维度是递增的,这样能够确保特征图在收缩部分和扩张部分相呼应,进而保证特征图能够恢复原尺寸。因为模型过多参数的更新和扩张部分的反卷积操作通常会造成数据的丢失,出现无法恢复原有尺寸的问题,这些问题增加了推理时的数据处理的难度。

1.2.2 扩张模块

扩张和收缩是对称的,这是典型的U形结构设计。扩张模块由3个结构相同的Transformer块组成,对应图2中的解码层。Transformer块的子模块堆叠顺序是相同,首先是归一化层,其次是多头自注意力层,第三层是第二个归一化层,最终是多层感知机。收缩部分传入的数据通过扩张部分第一个Transformer块的归一化层后,特征的差距减小,这使得传入多头自注意力层的数据能够在分头后每一部分都具有近似的数学性质,再做一次归一化依然可以起到这种作用。另外,还能够加快参数更新速度,进而减小模型的开销,并且保证模型的整体分割能力不会下降。每一个Transformer块输出特征的尺寸是递增的,维度是递减的,这样能够确保特征图在收缩部分传来的特征可以轻松地与扩张部分生成的特征合并,减少信息损失,提高了模型对数据的敏感度,有利于模型学习数据特征。

1.2.3 跳转连接

跳转连接的操作在收缩扩张这类U形结构中很常见,本文不仅使用了特征跳转连接还使用了注意力跳转操作。特征跳转连接就是将收缩部分的特征直接传入扩张部分,直接进行数学意义上的相加,相加后的新的特征作为扩张部分的特征,这种操作是具有开创性意义的,经过验证是有效的,模型的分割效果也是有明显提升的。Transformer在训练过程中会生成大量的数据,这些数据的计算很考验CPU的能力。为了解决这个问题,本文提出的模型将收缩部分计算出的K、V和Q值直接代替扩张部分Transformer对应层的对应值,这种操作使模型在计算上花费的时间较小,从而提升模型的训练速度。

2 实验

2.1 评价指标

为了更好地衡量模型的性能,采用了图像分割领域常用的评估指标Dice相似系数(dice similarity coefficient,DSC)来测量增强肿瘤区域(ET)、肿瘤核心区域(TC)以及整个肿瘤区域(WT)的分割精度。在脑胶质瘤图像分割中,Dice得分可以用来比较自动分割结果和人工分割结果(或真实标注)之间的一致性,反映分割的精度和覆盖度。豪斯多夫距离(hausdorff distance,HD)是一种用于衡量两个集合之间的距离的指标,常用于评价图像分割的精细度。

(6)

式中:TP为模型正确预测组织正样本的数量;FP为模型错误预测组织正样本的数量;FN为模型错误预测组织负样本的数量。

2.2 实现细节

本文所有实验都是基于Ubuntu16.04平台进行,模型的训练框架为Pytorch1.8,Python3.9。在训练模型时,将输入图像的尺寸缩放到128×128×128,学习率设置为1e-4,学习率调整策略设置为余弦衰减,优化器设置为Adam优化器,batch大小设为4,训练100个epoch,并采用包含32G内存的Tesla V100显卡对模型训练进行加速。网络的完整训练过程见图4,可以发现,随着迭代次数的增加,模型的训练损失稳步下降,直至收敛。

2.3 实验结果与分析

为了验证本文提出网络结构的有效性,在BraTS2021数据集上进行了大量的对比实验。实验结果见表1,本文提出的网络结构在ET、TC和WT区域上的Dice得分分别达到了83.51%、87.66%以及91.39%。相比其他5种分割模型,其中ET区域的得分比最高的3D U-Net提升0.12%,比最低的UNTER提升3.73%;TC区域的得分比最高的nnFormer提升1.18%,比最低的UNTER提升4%;WT区域的得分比最高的nnFormer提升1.02%,比最低的TransBTS提升2.14%。在性能指标上,本文提出模型的参数量只比3D U-Net高出8.9 M,但是浮点运算量却减少了392.9 GFlops。而相比TransBTS、UNTER、V-Net和nnFormer等模型,本模型的参数量分别减小了12.2、81.7、48.5和18.9 M。此外,本模型的浮点运算次数只比nnFormer增加了54.3 G,参数量却减小了近一倍。而相比TransBTS、UNTER、V-Net和3D U-Net模型,浮点运算次数则分别减少了168.0、28.5、600.9和392.9 G。综合整个实验结果,本文提出的分割模型在精度和速度上取得了较好的平衡,更接近临床应用的目标。

表1 实验性能对比

2.4 消融实验

为了探讨网络设计中重要参数对模型性能的影响,本节针对不同参数进行了消融实验。关于预训练权重对模型性能的影响结果如表2所示,在使用预训练权重后模型的整体性能都有所提升。其中,ET、TC、WT的Dice得分分别提高2.02%、1.34%、0.29%,可以发现,使用预训练权重对于难以识别的小目标提升更大,例如ET。

表2 预训练权值对模型的影响

关于跳转注意力对模型性能的影响结果如表3所示,在使用跳转注意力后,模型对ET和TC的Dice得分分别提高1.11%和1.39%,而WT的Dice得分仅提升0.13%。

表3 跳转注意力对模型的影响

3 结论

本文针对卷积神经网络在提取图像特征时存在的全局归纳偏差能力受限的问题,提出了一种完全基于Transformer的端到端的U形结构网络,能够直接处理三维的医学图像,并通过跳转连接降低了Transformer在小数据集上的过拟合风险。实验结果表明,本文提出的网络结构相比TransBTS、UNTER、V-Net和3D U-Net等网络具有更强的分割性能。

然而,本文在训练模型时,使用了Transformer在imagenet-1k数据集上的预训练权重对模型参数进行初始化,模型的分割精度有所提升,这说明Transformer模型对于大数据集的依赖程度比较高。如果不使用跳转注意力,不将模型收缩部分的K、V、Q值传入扩张部分,那么模型不仅会降低训练速度,分割性能也会降低。在下一步的研究工作当中,将针对这些问题对模型做进一步的优化,提高模型的识别精度与速度,争取达到临床应用的目标。