作者引入了D-Net用于体积分割医学图像,通过将动态大核模块和动态特征融合模块整合到分层 Transformer 架构中。

分层 Transformer 在医学图像分割中已取得了显著的成功,这归功于它们的大接收域以及有效利用全局长距离上下文信息的能力。卷积神经网络(CNNs)也可以通过使用大核来获得大的接收域,这使得它们能够在较少的模型参数下达到具有竞争力的性能。

然而,融入了大卷积核的CNN仍然受限于自适应地捕捉形状和大小变化较大的器官的多尺度特征,这是由于它们采用了固定大小的核。此外,它们也无法高效地利用全局上下文信息。为了解决这些限制,作者提出了动态大核(DLK)和动态特征融合(DFF)模块。DLK模块使用多个具有不同核大小和膨胀率的的大核来捕捉多尺度特征。随后,使用动态选择机制根据全局信息自适应地强调最重要的空间特征。

此外,还提出了DFF模块,以根据它们的全局信息自适应地融合多尺度局部特征图。作者将DLK和DFF集成在分层 Transformer 架构中,以开发出一种新颖的架构,称为D-Net。D-Net能够有效地利用多尺度大接收域并自适应地利用全局上下文信息。广泛的实验结果表明,D-Net在两个体积分割任务中的表现优于其他最先进的模型,包括腹部多器官分割和多模态脑肿瘤分割。

代码:https://github.com/sotiraslab/DLK

1 Introduction

视觉 Transformer (ViTs)的发展在计算机视觉任务上带来了显著的改进[8]。ViTs成功的关键因素是注意力机制,这使得基于ViT的模型拥有大的感受野,能够利用全局上下文信息贯穿整个输入图像。然而,由于在处理高分辨率图像时自注意力的计算复杂性高,ViTs在作为通用 Backbone 网络上面临着挑战。为了降低ViTs的复杂性,已经提出了分层ViTs。它们在建模不同尺度上的密集特征时更为高效,用线性复杂度近似自注意力。由于其卓越的性能,分层ViTs最近被用作医学图像分割的 Backbone 网络。然而,注意力机制常常限制了(分层)基于ViT的模型在有效提取局部上下文信息方面的能力。

另一种广泛使用的主干网络,卷积神经网络(CNN),在局部特征提取方面具有优势。然而,CNN的感受野受限于较小的卷积核。为了扩大它们感受野,引入了大的卷积核(LCK)并将其整合到CNN架构中。目前,基于LCK的CNN在医学图像分割中受到了关注。然而,这些网络依赖于单一固定大小的大核进行特征提取,这限制了它们捕捉具有大器官间和受试者间在形状和大小上变异的多尺度特征的能力。此外,它们缺乏增强局部特征与全局上下文信息之间交互的机制。

为了解决这些限制,作者提出了动态大核(DLK)和动态特征融合(DFF)模块。在DLK中,作者建议使用多个不同大小的深度卷积大核。这些核使得网络能够捕捉多尺度的上下文信息,有效地处理在形状和大小上的较大变化。与Atrous Spatial Pyramid Pooling (ASPP) 或其他并行设计中并行聚合这些核的方式不同,作者顺序地聚合多个大核以扩大感受野。随后,基于动态机制的思想,作者引入了一种空间上的动态选择机制,以根据全局上下文信息自适应地选择最有信息量的局部特征。

此外,DFF模块被采用以基于全局信息自适应地融合多尺度特征。在融合过程中,使用了通道上的动态选择机制来保留重要的特征图,然后使用空间上的动态选择机制来突出重要的空间区域。作者将提出的DLK和DFF模块集成到一个分层 Transformer 架构中,称为D-Net,用于3D体积分割医学图像。作者在两个分割任务上评估了D-Net:腹部多器官分割和脑肿瘤分割。提出的模型优于 Baseline 模型。

作者的主要贡献有三个:

  1. 作者提出了一种用于通用特征提取的动态大核模块。DLK采用多个大型卷积核来捕捉多尺度特征。随后,它利用动态选择机制,根据全局上下文信息自适应地突出最重要的空间特征。
  2. 作者提出了一种用于自适应特征融合的动态特征融合模块。DFF通过动态选择机制,根据全局信息自适应地融合多尺度局部特征。
  3. 作者提出了一个用于3D体积分割的D-Net。D-Net通过将DLK和DFF模块整合到分层ViT架构中,采用分层转换行为,以较低的模型复杂度实现了卓越的分割精度。

2 Method

Dynamic Large Kernel (DLK)

DLK. 作者提出了动态大核(DLK)方法,通过大感受野自适应地利用空间上下文信息(图1)。具体来说,作者使用多个大型深度方向核来提取多尺度特征。

D-Net_自适应

D-Net_自适应_02

2.1.2 DLK module.

D-Net_人工智能_03

2.1.3 DLK block.

为了利用分层Vision Transformers (ViTs) 的缩放能力,DLK块是通过将标准分层ViT中的多头自注意力替换为所提出的DLK模块来构建的。生成的DLK块包括一个DLK模块和一个MLP模块。

D-Net_人工智能_04

Dynamic Feature Fusion (DFF)

D-Net_人工智能_05

D-Net_人工智能_06

D-Net Architecture

D-Net的整体架构包括一个编码器、一个瓶颈层、一个解码器以及一个显著性层(图3)。显著性层用于从原始图像中提取显著的空间特征,而编码器-解码器架构负责学习层次化的特征表示。

D-Net_人工智能_07

3.2.2 Encoder.

D-Net_自适应_08

3.2.3 Bottleneck.

D-Net_自适应_09

3.2.4 Decoder.

D-Net_卷积核_10

2.3.5 Salience layer.

D-Net_人工智能_11

3 Experiments and results

3.0.1 Datasets.

D-Net_卷积核_12

3.0.2 Implementation details.

D-Net是使用PyTorch5实现的。损失函数采用了骰子损失和交叉熵损失的组合。在腹部多器官分割中,使用了AdamW作为优化器。初始学习率设置为0.0001,并应用了学习率衰减策略(ReduceLROnPlateau)。 

对于脑肿瘤分割,作者遵循了nnUNet 中的协议。使用SGD作为优化器。初始学习率设置为0.001,并使用多项式学习率调度器进行衰减。为了公平比较,所有实验都采用相同的设置并由作者实施。

3.0.3 Main results.

作者比较了D-Net与最近的几种最先进的分割模型在两个分割任务上的性能,包括3D U-Net (nnUNet) ,TransUNet,TransBTS,UNETR,nnFormer和3D UX-Net。

D-Net_3D_13

表1展示了在AMOS腹部多器官分割任务上的性能对比。D-Net以相对较少的FLOPs和最低的参数数量取得了最佳的整体性能。此外,D-Net在所有特定器官分割任务中的Dice得分都有显著提升。

D-Net_3D_14

表2展示了MSD脑肿瘤分割任务的结果。与其它分割方法相比,D-Net在所有分割任务上都展示了优越的性能。

D-Net_自适应_15

3.3.2 Ablation study.

D-Net_卷积核_16

4 Conclusion

作者引入了D-Net用于体积分割医学图像,通过将动态大核模块和动态特征融合模块整合到分层 Transformer 架构中。动态大核块被采纳为基本块,用于通用多尺度局部特征提取和自适应的全局空间信息利用。此外,还提出了动态特征融合模块以实现自适应特征融合。在两个分割任务上,即腹部多器官分割和脑肿瘤分割,D-Net的表现优于当前流行的 Baseline 。作者相信D-Net有潜力在各类医学图像分割任务上实现令人期待的分割性能。