集成多分枝残差块与融合注意力, MBDRes-U-Net 在多模态图像分割中实现高效特征提取与精确分割!

大模型机器学习算法

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

准确分割脑肿瘤在脑肿瘤疾病的诊断和治疗中起着关键作用,它是一种关键技术,用于量化肿瘤并提取其特征。随着深度学习方法的广泛应用,计算负担日益加重。

为了实现一个轻量级的模型并具有良好的分割性能,本研究提出了一个名为MBDRes-U-Net的模型,该模型使用三维(3D)U-Net编码器框架,集成了多分枝残差块和融合注意力。

通过分支策略,模型计算负担得以降低,有效地利用多模态图像中的丰富局部特征并增强子肿瘤区域的分割性能。

此外,在编码过程中,引入了一种自适应加权扩展卷积层到多分枝残差块中,从而丰富了特征表达并提高了模型的分割准确性。在2018年和2019年的Brain Tumor Segmentation(BraTS)挑战赛数据集上的实验表明,该架构可以在保持脑肿瘤分割高精度的同时,显著降低计算开销。

作者的代码https://github.com/Huaibei-normal-university-cv-laboratory/mbdresunet。

1 Introduction

恶性脑肿瘤中最常见的是胶质瘤,可分为低度胶质瘤(LGG)和高度胶质瘤(HGG)两种亚型。胶质瘤具有高发病率、高复发率和高死亡率,但治愈率较低,治疗具有挑战性。治疗还需要准确地进行医学影像检查和图像处理与分析,这主要依赖于医生。因此,准确分割脑肿瘤是医学诊断和病理分析的关键,如图1(e)所示。这一任务包括划分几个子区域,即增强核心(EnC)、肿瘤周围水肿(PTE)和非增强核心(NEC)区域[1]。在许多医学影像方法中,磁共振成像(MRI)在治疗胶质瘤方面具有显著优势。它可以提供丰富的肿瘤信息,且是非侵入性的;此外,在成像过程中不会暴露患者身体受到辐射。有四种常见的成像模式(图1(a)-图1(d))--即液体衰减反转恢复(FLAIR)、T1加权(T1)、对比增强T1加权(T1c)和T2加权(T2)模式。值得注意的是,不同的成像效果强调不同的组织特性和肿瘤扩散区域[2]。医生通常手动 Token 大量MRI扫描,逐层逐块;肿瘤区域的划分取决于他们的经验和专业知识。这项工作既费时又费力,意见分歧较大,因此准确自动分割具有重要的研究意义。

picture.image

许多尝试已经用于解决多模态MRI脑肿瘤分割问题,包括边缘分割,基于区域的[3],基于图谱的[4],以及基于机器学习的方法。尽管它们具有较低的计算复杂性,但这些方法严重依赖人工标注和大量数据集进行预训练,这可能耗时且导致分割性能不佳。

自从2012年以来,多模态脑肿瘤分割(BraTS)挑战(由国际医学影像计算与计算机辅助干预协会(MICCAI)及其他组织联合赞助)凸显了深度学习算法在脑肿瘤分割任务中的进步。U-Net及其变体在多模态脑肿瘤分割任务中的性能取得了显著的进展。从三维(3D)U-Net模型[6]到TransBTS[7]和SwinUNETR模型[8],研究不断提高了脑肿瘤分割方法的表现。此外,从最初将U-Net架构集成到各种尝试引入 Transformer 的模型,这些模型为脑肿瘤分割领域带来了创新和进步。NN-U-Net[9]确认了U-Net架构在图像分割任务中发挥的重要作用,基本U-Net架构在多个BraTS任务中都是有效和有竞争力的。

此外,由于注意力机制使神经模型能够准确关注输入信息的所有相关元素,它已成为提高深度神经模型性能的重要组件。例如,SENet[10](在2017年的ImageNet竞赛中获胜)中的挤压与激励(SE)模块从通道的角度使用权重矩阵为图像的不同位置分配不同的权重,以获取更重要的通道特征信息。然而,与仅关注通道的注意力机制相比,将通道和空间注意力机制结合的模块可以实现更好的效果。例如,卷积块注意力模块(CBAM)[11]依次推理出沿通道和空间维度的注意力图。

尽管3D卷积神经网络(CNNs)在脑肿瘤分割方面取得了显著成果,但这些模型通常过于复杂,计算和生成大量参数,导致计算开销高。为了解决这个问题,已经尝试使用轻量级模型架构;然而,轻量级模型的分割性能仍无法与高级模型相比。在保持较低计算成本的同时提高分割精度仍是一个难以解决的问题。

因此,作者提出了一种基于多模态的3D脑肿瘤轻量级模型(MBDRes-U-net),该模型不仅解决了基于单模态图像表示的单一特征问题,还减轻了复杂模型的计算负担。在所提出的模型中,使用了基于多分枝并行卷积的残差块来替换常见的3D卷积块;通过使用组卷积来降低计算复杂性。在编码器中引入了一种自适应3D膨胀卷积操作,以获得多尺度特征表示,并在编码阶段添加了一个3D多注意力模块(SCA3D),以使CNN能够关注肿瘤区域。所提出的模型在参数数量和模型复杂性方面都具有其他同类型子模型的特定优势。作者在BraTS 2018和BraTS 2019数据集上评估了该模型,在两个数据集上都获得了良好的分割性能。本研究的贡献可以总结如下:

在本研究中,作者提出了一种3D轻量级编码-解码模型MBDRes-U-net,用于脑肿瘤的多模态磁共振图像分割。MBDRes-U-net是基于对称编码-解码结构,集成了新的残差块和注意力机制。由于其简单的模型结构,它可以作为3D脑肿瘤分割的基准,从而推动脑肿瘤MRI图像分割的研究。

提出了一种基于多分支空洞卷积的新残差块,同时扩大了接受域并减少了参数数量,从而解决了由3D卷积引起的计算开销问题。此外,作者引入了自适应空洞卷积而非普通空洞卷积,以丰富特征表示。

提出了一种可插拔的3D注意力模块(3D SACA),更适合脑肿瘤分割。探讨了空间和通道注意力之间的改进结合机制。

2 Related Work

Convolutional Neural Network Model

近年来,深度学习方法,尤其是卷积神经网络(CNNs)在医学图像处理应用中得到了广泛应用。CNN-based模型可以通过学习准确捕捉二维(2D)和三维(3D)医学图像的局部特征。Kamnitsas等人提出了一个全连接的多尺度模型架构(DeepMedic)[12]来分割3D脑肿瘤图像。Chen等人提出使用CNN的密集连接来提高脑肿瘤分割的准确性[13]。随着U-Net的出现[5],基于 Short-Cut 的编码器-解码器对称架构(在医学图像分割中的优势)越来越明显,近年来在脑肿瘤分割中得到了广泛应用。然而,使用医学图像的2D切片所引起的数据损失很难忽视。为了更好地学习3D数据的成像特征并满足提供医疗援助的临床准确性要求,Cicek等人将U-Net模型从2D扩展到3D图像[14]。随后,Wang等人提出了一种基于3DU-Net模型的脑肿瘤分割模型[15]。最近,Jiang等人提出了一种两阶段的级联U-Net模型,在2019年BraTS 2019 Challenge Segment Task中获得第一名的成绩[16]。TransBTS [7]用 Transformer 集成块替换了3DU-Net [14]的 Bottleneck 层,以提取更多的全文本信息并补偿早期模型建立长连接的能力不足。来自德国癌症研究中心的参赛团队提出了一种基于nnU-Net模型的脑肿瘤分割方法,在2020年BraTS挑战分割任务中获胜[9]。Hatamizadeh等人提出了一种基于U-net架构的全 Transformer 编码器模型,即SwinUNETR [8]模型,对脑肿瘤分割产生了重大影响。2021年的获奖模型是基于nnU-Net模型改进的[17]。

随着CNN技术的快速发展,模型的计算成本得到了显著提高。考虑到基于CNN的模型的效率,许多轻量级用于脑肿瘤分割的模型已经提出。Chen等人使用可分离的3D卷积(S3D-U-Net)来降低计算成本和内存需求[18];然而,对于增强肿瘤区域的分割精度,尤其是低通量。Zhou等人[19]提出了一种使用轻量级ShuffleNetV2模型[20]作为编码器,引入残差模块解码器以避免降质的3D残差神经模型(ERV-NET)。Chen等人提出了一种新的3D扩展多纤维模型(DMF-Net),其中分割的多尺度图像表示通过引入加权3D扩展卷积操作获得,从而减少了模型参数的数量并实现了准确的分割[21]。然而,尚未解决通道信息交换的问题。HMNet模型[22]使用并行的高分辨率多级分支提取多分辨率特征,进一步降低了模型的复杂性和计算开销。

2.2 Attention Mechanism

为了进一步提高肿瘤分割模型的准确性,可以将注意力机制引入以聚焦于肿瘤相关区域。在医学图像处理中最常用的注意力机制包括通道和空间注意力,它们通过使用不同的聚合策略、变换和增强函数,将所有位置的同类型特征聚合起来,从而增强原始特征。例如,Mobarakol等人采用了结合通道和空间注意力的3D U-Net架构,用于分割[23]。MBANet模型[24]包括3D多分枝注意力,其中3D空间注意力(SA)作为注意力层在编码器中提供通道和空间注意力。受到SA的启发,3D SACA注意力模块使用通道切乱[20]来促进通道间的信息 Stream ,而不生成参数。然而,初始分组操作被丢弃,3D SE被用作3D通道激励模块和具有3D卷积的空间激励模块。作为即插即用的注意力模块,3D SACA模块更适合处理3D卷积模型的3D图像。

3 Methods

MBDRes-U-Net

图2展示了MBDRes-U-Net模型的架构。该算法使用3D U-Net结构编码器框架。考虑到分割目标肿瘤区域位于MRI输出中,作者在编码器前面添加了一个多分枝3D SACA混合注意力模块,以便在提取特征之前使模型更关注感兴趣的区域。

picture.image

在编码阶段,首先使用了一个卷积核,步长为2进行下采样,以减少显示内存。然后通过六个多分支扩展卷积残差(MBDRes)块和自适应扩展卷积层提取多尺度特征。在解码器中,使用三线性插值方法对特征图进行上采样。这些上采样的特征图然后与来自跳过的连接的高分辨率特征进行拼接。接下来,拼接的特征经过解码卷积块和MBRes块逐步恢复原始分辨率。然后,使用步长为1的卷积将所有通道信息进行融合。最后,使用SoftMax函数通过端到端分割实现分割图。

3D CNN时代的到来使得MRI 3D数据的特性得到了充分利用。然而,当3D卷积核在全部特征映射通道上运行时,计算复杂性(即每秒浮点运算次数,FLOPS)呈指数增长。因此,在训练过程中3D CNN具有较高的计算成本。

在ResNet v2 [25]中使用的预激活残差块(图2(a))扩展为3D卷积模型(图2(b)),可以引入分组卷积--一种有效的模型加速方法(包括ResNeXt [26]和ShuffleNet [27]模型)--来缓解计算负担。假设ResNet v2单元[25]被分为个并行分支, Kernel 大小保持恒定。

参数数量在多分支分组后的残差块为:

该值减少了倍。

在本工作中,、 和 分别表示自适应扩张卷积层中每个分支的权重, 表示扩张率。这种分组策略可以有效减少模型参数数量并加速计算。然而,由通道分组生成的多个分支独立并并行工作,影响通道间的正常信息交互,降低模型的学习能力。为解决缺少通道信息交换(图2(c))的问题,作者可以在模型中添加一个通道交换模块。

1×1×1卷积层在残差块的开始和结束部分,用于在每条支路之间传递信息。此外,残差连接可以布置在单元外部,这样低层的信息可以直接传递到高层,而无需生成额外的参数,从而提高模型的学习能力。因此,MBR块可以表示为MBRes块。

考虑到传统卷积核在卷积中的大小有限(导致其接受域有限),为了扩大模型的感受野,学习脑肿瘤MRI的多尺度特征,并捕捉3D空间相关性,作者在编码部分引入了一种自适应加权扩展卷积层来替代传统的卷积操作,从而得到了多支扩展卷积残差块(称为MBDRes块)。捕捉多尺度信息是一种有效的策略,以前已经成功地使用过。Tokunaga等人提出了一种在病理学中使用三个并行CNN和加权 ConCat 来提取多尺度信息的语义分割任务[28]。在Chen等人提出的密集融合最大池网络(DFMN)[21]中,使用了三个具有不同扩展率的并行卷积层作为加权和。

适应性加权扩展卷积层的结构包括三个并行的3D扩展体积积分分支,每个分支的扩展率分别为1、2和3(见图3(d))。在每个分支初始化后,为其分配三个权重(和),然后将每个分支的结果相加。这种权重的初始化确保了初始阶段每个分支对模型具有相同的影响。

picture.image

3D Saca

存在显著不平衡的BraTS数据集,其中肿瘤区域仅占MRI图像的1.5%,增强型肿瘤(ETs)仅占整个肿瘤(WT)图像的11% [29]。为了消除大面积背景对分割的影响,平衡了模型在空间细节和高级形态特征学习能力,使模型更加关注肿瘤区域。

可以将3D注意力机制引入,通过3D空间和通道之间的特征关系提取3D空间和通道注意力,如下所示:

picture.image

其结构如图4所示。可以引入一种通道分裂[20]操作来将输入特征X(X∈R^{1/2c×h×w×d})通过通道分裂成两部分,分别记作Xs(Xs∈R^{1/2c×h×w×d})和Xc(Xc∈R^{1/2c×1×1×1})。其中,C、H、W和D分别表示特征图中的通道数、高度、宽度和深度。Xc首先进行平均池化,以获得全局通道信息,然后将信息传递到通道激励模块,以获得i/2c×1×1×1的通道相关性。Xs被引入到空间激励模块,通过3D卷积层的点乘操作将空间特征相关性聚合到维度,从而实现空间注意力权重。

picture.image

因此,该模型能够自适应地调整不同空间位置的特征响应。

接下来, 和 进行聚合。然后,作者将残差进行融合以减少并行激发引起的稀疏性。最后,使用通道重新排列解决分支策略引起的信息交换问题。

4 Results and Discussion

Experimental Details

4.1.1 Dataset and data pretreatment

作者在以下数据集上评估了MBDRES-U-NET模型:

BraTS 2018 数据集:训练数据集包含 285 个样本;验证数据集包含 66 个样本。

2019年BraTS数据集:训练数据集包含335个样本;验证数据集包含125个样本。

每个样本的体积为240×240×155。用于肿瘤分割的标签包括背景(标签0)、坏死和非增强型肿瘤(标签1)、肿瘤周围水肿(标签2)以及GD增强型肿瘤(标签4)。为了便于模型训练,数据集进行了如下预处理:

由于数据集通过多种途径获得,导致强度分布不均;因此,作者使用Z-score方法对MRI图像进行标准化。

脑肿瘤图像的背景信息对分割没有意义,因此作者随机裁剪了数据为128×128×128 voxel输入。

为了防止过拟合,作者采用了以下数据增强策略来增加训练数据:

0.5 对于轴向、冠状和矢状随机反转;随机旋转角度区间为[-10°, , , , , 10°+10°]。

4.1.2 Assessment indicators

该模型的有效性基于计算复杂性和分割精度进行评估。通过提供的一个在线门户,使用BraTS挑战赛的验证数据集进行验证。具体而言,分割精度通过Dice系数和Hausdorff距离(95%)进行测量,其中ET、WT和TC分别表示增强肿瘤区域(标签1)、整个肿瘤区域(标签1、2和4)和肿瘤核心区域(标签1和4)。复杂性可以通过 Params 和 FLOPS 指标确定。 Params表示模型的空间复杂性,FLOPS 表示其时间复杂性,如下所示:

picture.image

其中,, 和 分别表示卷积核的高度、宽度和深度, 和 分别表示输入和输出通道的数量,而参数 h,w 和 分别表示图像的高度、宽度和深度。

4.1.3 Experimental setup

作者在Python 3.6中运行了实验代码,使用了16个批次,在3个并行NVIDIA A30 GPU上训练模型500个周期,并使用PyTorch框架构建了所有实验模型。使用了Adam优化器,学习率设置为0.001。

4.2 Comparison Experiments with State-of-the-art Methods

为了验证所提出模型的性能,作者将MBDRes-U-Net的分割性能与其他先进模型(包括3D U-Net、基于U-Net的CNN与Transformer以及其他轻量级模型)在BraTS 2018和2019数据集上进行了比较。比较结果如表1和表2所示。与非轻量级模型相比,所提出的模型在模型复杂性方面具有优势,从而实现了对WT和TC更高效的分割。

picture.image

picture.image

MBDRes-U-Net模型的参数是传统3D U-Net模型的四分之一(表1)。此外,计算复杂度降低了1643.75 G,分割精度显著提高(分别比ET、WT和TC段的3D U-Net模型提高了3.2%、1.8%和13.6%)。与3D-ESP-Net和S3D-U-Net模型相比,尽管MBDRes-U-Net模型的参数数量略有增加,但计算复杂度仅为后者的三分之一,分割精度也有显著提高。与具有相似参数的DMF-Net模型相比,虽然MBDRes-U-Net模型的Dice系数降低了0.1%,但WT和TC段的Dice系数分别提高了0.7%和1.8%。尽管所提出的模型不如HMNet模型轻量,但ET、WT和TC段的Dice系数分别提高了0.5%、0.2%和1.0%,Hausdorff距离分别减少了0.002、0.389和2.037 mm。总体平均分数比HDC-Net模型提高了0.6%,计算复杂度接近。与最新的ADHDC-Net脑肿瘤分割模型相比,虽然MBRes-U-Net模型在ET段的分数降低了0.3%,但WT和TC段的分数分别提高了0.5%和0.3%,MBRes-U-Net模型的平均分数比非轻量级ADHDC-Net模型提高了0.17%。因此,所提出的方法是一种更高效的算法,可以实现与轻量级HMNet模型相当的分割精度。

表2列出了使用BraTS 2019数据集的MBRes-U-Net模型的结果。通过重新训练模型,MBRes-U-Net模型比3D U-Net模型更轻便且高效。它比3D ESP-Net和DMF-Net模型的参数更少,分割精度更高。与最新的脑肿瘤分割模型相比,MBRes-U-Net模型在ET和CT方面的改进显著,分别比MBANet模型提高了0.09%和0.5%,比ADHDC-Net模型提高了0.8%和1.3%。尽管MBRes-U-Net的ET和WT Dice分数降低了0.5%,但TC Dice_score比TransBTS模型提高了1.6%。与具有相似计算复杂度的HDC-Net模型相比,TC优势明显(增加了2.6%),整体平均值增加了1.03%。与HMNet模型相比,尽管所提出的模型参数比HMNet模型多3M,但WT Dice_score降低了0.4%,而ET和CT的TC分别提高了1.1%和0.5%,计算复杂度减少了103 G。因此,可以看出MBRes-U-Net模型在轻量级和非轻量级模型中更具竞争力。

此外,Fig 5中展示了MBDRes-U-Net模型的多种可视化结果。不同颜色代表不同类型的肿瘤,即红色区域表示坏死和非增强性肿瘤,黄色区域表示增强性肿瘤,绿色区域表示水肿。

此外,从左到右,将FLAIR、DMF-Net、HDC-Net、ADHDC-Net和MBDRes-U-Net模型的分割结果叠加在FLAIR图像上。如所示,SGEResU-Net模型能够有效地分割增强肿瘤、总体肿瘤和核心肿瘤区域。

Ablation Experiment

4.3.1 Adaptive weighted expansion convolution layer

比较消融实验如表3中所示,用于验证是否需要自适应加权扩张卷积层和自适应加权算法。

picture.image

与没有自适应加权扩展卷积层的设计方案相比,显然,扩展卷积可以提高Dice分数。加权策略的有效性通过与等权重方案()进行比较得到证明。由于该权重策略能够学习和自适应地选择多尺度上下文信息,因此,这种加权策略导致了更优的分数,尤其是在WT指标上。

训练过程中使用的权重,和在图6中显示。显然,每个单元的权重参数都处于收敛状态,反映了多分支扩展卷积残差块的作用。值得注意的是,(蓝色线,对应小感受野)在第一、二、三块中起着重要作用,而在更高层次的块中被削弱。然而,显然,(红色线,对应大感受野)的扩展分支在MBDRes Block-2-6中起着主导作用,这可能是因为具有较小感受野的核无法在较小的维数下捕获更高层次的有用语义信息。

picture.image

MBDResk块,以及MBDRes块-2-6等。图片中的蓝色线表示,橙色线表示,绿色线表示。

4.3.2 Multibranch fused attention

针对BraTS 2018数据集进行了消融实验,以评估注意力模块和并行分支策略的必要性以及有效性。

参考文献

[0]. MBDRes-U-Net: Multi-Scale Lightweight Brain Tumor Segmentation Network.

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论