Trans-SegNet:一种基于Transformer的脑肿瘤图像分割网络
2024-01-04仇龙
仇龙
(北京科技大学,北京 100083)
脑肿瘤是一种常见的神经系统疾病,是指在脑部或其周围组织中发生的一类恶性或良性肿瘤,其分布范围广泛,易于侵袭周围组织,对患者的生命和健康造成威胁。全球每年新发脑肿瘤约25万例。据统计,脑肿瘤在恶性肿瘤中占比不到2%,但是近年来其发病率持续上升,已成为21世纪危害人类健康的重要疾病之一。早期诊断和治疗是脑肿瘤预防和根治的基础[1],但医生在诊断中会受到自身经验、知识水平、情感因素等的影响,导致诊断结果存在主观性,容易出现误诊和漏诊等问题。此外,由于医生的不同背景、经验和知识水平的不同,不同医生对同一病例的诊断结果可能缺乏一致性。而医学图像分割技术是一种基于计算机算法的自动化技术,可以通过降低医生分析的主观因素,提高疾病诊断的客观性和准确性。还可以帮助医生更准确地评估肿瘤的位置、形状和大小等信息,通过对图像进行分割,可以清晰地显示出肿瘤和周围组织之间的分界线,减少误诊和漏诊的风险。对于需要手术治疗的患者,医学图像分割技术可以帮助医生更准确地评估肿瘤的位置和大小,规划手术方案,减少手术风险。
随着深度学习技术的不断发展,医学图像分割中基于卷积神经网络的方法已经成为主流[2-3],其中,UNet[4]、V-Net[5]和U-Net++等方法得到广泛应用。UNet 基于卷积神经网络采用了编码器-解码器结构。在编码器部分,卷积和池化操作可用于降低输入图像的分辨率并提取图像特征。在解码器部分,卷积和上采样操作被用来将特征图恢复到原始图像的分辨率,并进行像素级别的分类。此外,还引入了跳跃连接,将编码器的特征图送入解码器中帮助解决分割任务中的信息丢失问题。虽然已经取得了较好的分割效果,但其由于卷积操作引入的固有偏置原因,仍存在无法建模长距离依赖和捕捉全局上下文信息的问题,这限制了分割精度和辅助诊断准确率的进一步提升[6-7]。近年来基于自注意力机制的神经网络结构Transformer[8]在自然语言处理领域和图像领域都取得了很好的效果,相比于传统的循环神经网络,Transformer 能够处理更长的文本序列,并且能够捕捉更复杂的语言关系。这使得Transformer 成了自然语言处理领域的一种重要技术。其通过内部的自注意力机制可以自动对图像中不同位置的特征进行加权和组合,从而提高模型的表达能力并建立全局层面的长距离依赖。这种特性在医学图像分割中非常重要,因为医学图像中存在很多微小的细节和不规则的形状,需要模型具有较强的表达能力才能更好地分割[9]。此外,全局信息可以帮助模型更好地理解图像中的结构和背景,在脑部MRI 图像中,肿瘤结构往往与大脑组织的结构紧密相关,需要考虑整个脑部结构才能更好地分割肿瘤。综上,本文从脑肿瘤图像分割算法中引入Transformer结构并利用自注意力机制,设计了新的医学图像分割网络Trans-SegNet。实验结果表明,Trans-SegNet可以有效地提升模型的建模能力并建立长距离全局信息,提升脑肿瘤分割的准确性。
1 Trans-SegNet分割算法设计
1.1 算法整体结构
本文基于传统的U-Net 架构和Transformer 结构,提出了一个新的脑肿瘤图像分割算法,其整体框架如图1 所示。相较于原本U-Net,本文将Transformer 层多尺度并行插入原始U-Net 的跳跃连接中并设计全局融合模块,补充不同尺度特征中的长距离依赖,并送入解码器进一步融合。具体来说,所设计模型Trans-SegNet 包含3个主要部分,分别是编码器、解码器以及全局特征融合模块。其中,编码器由5个卷积块组成,每个块由3×3卷积、批归一化和ReLU激活函数组成,其结构如图1所示。并在每个卷积块间逐步下采样压缩图像尺寸,依次为原始图像分辨率的1/4、1/8、1/16和1/32。同时,加入残差连接的操作,使模型优化更快避免出现训练不稳定的现象。得到多尺度特征后,由于Transformer 是处理序列信息的,需要首先对特征图进行特征形状变换,将3维的图像特征变换为1维。为保留图像块之间的相对位置信息,在序列上附加位置编码后,分别送入4个并行的全局特征融合模块来同时建立4个尺度上的长距离特征关系。挖掘并增强肿瘤相似特征间的相关性,提取上下文语义信息,帮助模型更好地理解图像中的结构和关系,区分肿瘤和正常组织,提升分割精度的准确性。将经过全局增强模块后的特征图进行形状变换,恢复其3维结构,使其适配3维卷积的形状结构后,分别送入解码器中。解码器中包含与编码器类似的卷积块,尺寸最小的特征图通过卷积模块和上采样操作与经过全局增强后的多尺度特征拼接融合并逐阶段恢复和生成用于最终分割的特征,实现脑肿瘤图像的精准分割,下面将详细介绍全局融合模块。
图1 Trans-SegNet模型结构图
1.2 全局特征融合模块结构
全局特征融合模块由自注意力机制、频域表示模块、层归一化和前向传播层等组件构成。图像送入全局特征融合模块前在一个图像块层面被序列化,将3维的医学图像数据转化成1维的序列数据,并附加位置编码保持图像相对位置关系。之后并行通过自注意力机制和频域表示学习模块建立整张图像上的相互作用并捕捉上下文信息。在自注意力机制中,对于每个图像块,会将该图像块的特征表示作为查询(Query)、键(Key)和值(Value),然后计算该图像块与所有其他该图像块之间的相似度。将所计算的相似度作为权重,对所有值进行加权求和,得到的加权和为该查询的输出[7]。其数学表达式如下:
式中,WQ,WK,WV表示可学习的参数,dk表示特征向量维度,x表示输入特征。通过上述先计算全局图像块的关联度,得到关联亲密度矩阵,再增强关联性强的图像块,抑制关联性弱的图像块的方法,模型的全局建模能力得到增强,图像特征之间的相互作用被构建。
为进一步完善和补充图像全局信息,本文设计了并行的频域表示学习模块。由于频域可以反映图像中的全局信息,其提供了对图像中所有频率的解析能力,可以通过分析频域中的频率,了解图像中存在的所有周期性变化。包括图像中的低频分量(代表图像中的大型特征和整体结构)和高频分量(代表图像中的细节和纹理)[9],因此在频域中分析是个合理可行的方案。具体来说,本文首先将特征经过傅里叶变换映射到频域中,之后通过添加2个可学习的参数增强模型在频域中丰富模型的全局特征表示能力,然后进行傅里叶反变换,与自注意力分支特征对齐,其数学表达式如下所示。
式中,FDB表示频域表示学习模块,IFFT、FFT分别表示傅里叶变换和反变换,a,b表示可学习参数,其形状大小与输入特征x相同,分别进行逐元素相乘和相加,通过这种方式可以有效地实现在频域中学习全局信息。
特征经过全局信息增强后被送入前向传播层中,其中包括两个全连接层和一个激活层,激活层引入了非线性映射,补足自注意力机制无法处理非线性关系的问题。全连接层可以进一步增强模型的表示能力。其数学表达式如下:
式中,W1,W2,b1,b2表示可学习的参数,此外,全局特征融合模块还引入了残差连接和层归一化等技术,来加速模型的训练和提高模型的泛化能力。
综上,全局特征融合模块通过自注意力机制和频域表示模块捕获增强了全局上下文信息,并从频域角度完善特征间相互依赖。此外,还加入了前向传播层以提升模型表示能力,实验证明,通过以上改进可以显著提升模型的分割精度。
2 实验与结果
2.1 实验细节和评价指标
本文基于PyTorch 框架并在2 张NVIDIA Geforce RTX3090 上训练,显存为24GB,使用PyTorch 版本为1.8.0,使用版本为11.3 的CUDA 和8.2 的cuDNN 进行训练加速。实验数据集为公开数据集BraTS2019[10],它包括335个用于训练的患者病例和125个用于验证的病例。每个样本都由具有4 种模式的3D 大脑MRI扫描组成。每个模态都有240×240×155的体积,已经对齐到同一空间,共包括背景、肿瘤全部区域(WT)、肿瘤核(TC)和增强肿瘤(ET)4 个类别。输入经过裁剪后的图片尺寸为128×128×128 以对齐空间分辨率和降低计算复杂程度。为避免过拟合,在实验中增加随机翻转、调节对比度和随机缩放以生成更多训练样本,并在测试时使用多尺度增强测试来提升模型性能。本文在通过五折交叉验证,评估得分训练样本数为269,验证样本数为66,批处理大小4,学习率1e-4,训练轮数500轮。本文以Dice得分来评估模型性能,其具体数学计算方式如下式所示。
式中A,B分别表示该类别图像上的预测区域和真实标签区域。此外,由于医学脑肿瘤图像存在前后景比例悬殊大,相对结构较为固定并且不同类别数量极度不平衡,本文不采用传统图像分割上的交叉熵损失,训练阶段优化的目标函数使用Dice 损失。Dice损失是基于相似性的评价指标,更看重重叠的比例。可以很好地处理类别像素值不平衡的问题。优化方面,由于传统随机梯度下降,可能会存在局部最优解和收敛速度慢的问题,本文选用了adam 的优化方法提升模型收敛速度。
2.2 实验结果
本文在BraTS2019 上进行实验,并验证了所设计模型和具体模块的有效性,实验结果如表1所示。在相同实验条件下,Trans-SegNet 相较于传统基于卷积神经网络的U-Net 模型在三种评价指标下均取得了大幅提升,其中ET 提升2.07%、WT 提升0.98%、TC 提升1.89%,证明了本文方法的优势。具体来说,从模型角度分析,为了排除是骨干网络的变化导致的精度差异,笔者设计了当Trans-SegNet 不添加全局特征模块仅采用卷积神经网络的实验,结果相较于之前方法并没有取得突出的性能优势。为了进一步地验证Transformer 插入到模型中的效果,笔者增加全局特征模块但不包含频域特征学习模块的实验,分割精度取得较大幅度提升1.31%(ET)、0.75%(WT)、0.98%(TC),实验结果证明,Transformer 的全局信息建模能力和多尺度插入融合方法,能够有效地帮助模型理解上下文语义,提升分割准确率。此外,为验证频域表示学习模块的有效性,笔者进行了添加频域表示学习模块的实验即最终的Trans-SegNet,结果表明通过频域表示学习模块,模型能在频域有效地完善全局信息,并进一步提升模型分割精度,实现脑肿瘤的精准分割。
表1 实验结果精度对比
3 结论
针对目前深度学习算法在脑肿瘤图像分割中存在全局特征和上下文信息提取能力弱以及分割精度低等问题,本文基于传统的卷积神经网络U-Net 算法提出了改进,本文通过在跳跃连接中添加全局特征增强模块,辅助模型捕捉全局语义信息。具体来说,其中包括由频域表示学习和自注意力机制组成的双分支全局信息增强部分和前向传播层以增强模型表示能力。实验结果表明,相比改进前的传统卷积神经网络U-Net分割精度分别在ET/WT/TC上提升了2.07%、0.98%、1.89%。并进一步从实验上证明了所设计频域表示学习模块和全局表示学习模块的有效性以及Transformer 和在跳跃连接中插入的方式能够有效地帮助模型具备全局建模能力,这有助于未来对脑肿瘤图像分割的研究,给医学图像辅助诊断技术提供借鉴。