点击下方卡片,关注 「AI视界引擎」 公众号
分层 Transformer 在医学图像分割中已取得了显著的成功,这归功于它们的大接收域以及有效利用全局长距离上下文信息的能力。卷积神经网络(CNNs)也可以通过使用大核来获得大的接收域,这使得它们能够在较少的模型参数下达到具有竞争力的性能。
然而,融入了大卷积核的CNN仍然受限于自适应地捕捉形状和大小变化较大的器官的多尺度特征,这是由于它们采用了固定大小的核。此外,它们也无法高效地利用全局上下文信息。为了解决这些限制,作者提出了动态大核(DLK)和动态特征融合(DFF)模块。DLK模块使用多个具有不同核大小和膨胀率的的大核来捕捉多尺度特征。随后,使用动态选择机制根据全局信息自适应地强调最重要的空间特征。
此外,还提出了DFF模块,以根据它们的全局信息自适应地融合多尺度局部特征图。作者将DLK和DFF集成在分层 Transformer 架构中,以开发出一种新颖的架构,称为D-Net。D-Net能够有效地利用多尺度大接收域并自适应地利用全局上下文信息。广泛的实验结果表明,D-Net在两个体积分割任务中的表现优于其他最先进的模型,包括腹部多器官分割和多模态脑肿瘤分割。
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 模型。
作者的主要贡献有三个:
- 作者提出了一种用于通用特征提取的动态大核模块。DLK采用多个大型卷积核来捕捉多尺度特征。随后,它利用动态选择机制,根据全局上下文信息自适应地突出最重要的空间特征。
- 作者提出了一种用于自适应特征融合的动态特征融合模块。DFF通过动态选择机制,根据全局信息自适应地融合多尺度局部特征。
- 作者提出了一个用于3D体积分割的D-Net。D-Net通过将DLK和DFF模块整合到分层ViT架构中,采用分层转换行为,以较低的模型复杂度实现了卓越的分割精度。
2 Method
Dynamic Large Kernel (DLK)
DLK. 作者提出了动态大核(DLK)方法,通过大感受野自适应地利用空间上下文信息(图1)。具体来说,作者使用多个大型深度方向核来提取多尺度特征。
此外,作者不平行结合多个核,而是将这些大型核以递增的核大小和膨胀率级联起来。这一设计有两个优点。首先,上下文信息在感受野内递归地聚合,使得有效的感受野能够逐步增大。其次,在更深和更大的感受野内提取的特征对输出的贡献更为显著,这使得DLK能够捕获更精细和更具有信息性的特征。
在作者的工作中,作者使用了两个带有大核的深度卷积(DWConv)层:,具有膨胀率为1的核,以及,具有膨胀率为3的核对第层的输入特征进行卷积:
通过级联这些核,DLK具有与核相同的有效感受野。通过沿着通道从级联特征应用平均池化(AVP)和最大池化(MAP),这些局部特征的全局空间关系被有效地建模。
然后使用一个 卷积层(Conv)来允许这些信息在不同的空间描述符之间进行交互,并使用Sigmoid激活函数来获得动态选择值 ,:
不同大核的特征通过利用这些选择值对它们进行校准,从而自适应地选择。最后,应用一个残差连接。
2.1.2 DLK module.
DLK模块是通过将DLK集成到两个线性层(卷积层;Conv)之间并加入GELU激活函数来构建的。还应用了残差连接。因此,DLK模块中第层的输出可以计算为:
2.1.3 DLK block.
为了利用分层Vision Transformers(ViTs)的缩放能力,DLK块是通过将标准分层ViT中的多头自注意力替换为所提出的DLK模块来构建的。生成的DLK块包括一个DLK模块和一个MLP模块。
与分层ViT块类似,在每个DLK模块和MLP模块之前应用了一个层归一化(LN)层,并且在每个模块之后应用了一个残差连接。因此,第层和第层中的两个连续DLK块可以计算为:
Dynamic Feature Fusion (DFF)
作者提出了一个动态特征融合(DFF)模块,用于根据全局信息自适应地融合多尺度局部特征(图2)。这是通过在融合过程中动态选择基于它们全局信息的重要特征来实现的。具体来说,特征图 和 沿通道维度进行拼接。
为了确保后续的块能够采用融合特征,需要一个通道减少机制将通道数恢复到原始数量。在DFF中,不是简单地使用一个 卷积来进行通道减少,而是通过全局通道信息 来引导。这一信息通过级联一个平均池化(AVGPool)、一个卷积层(Conv)和一个Sigmoid激活来描述特征的重要性。
融合特征由全局通道信息进行校准。随后,使用一个 卷积层(Conv)根据特征的重要性来选择特征图。此通道信息将指导卷积层在保留重要特征的同时,丢弃信息量较少的特征。
为了模拟局部特征图之间的空间上的相互依赖关系,全局空间信息 通过 卷积层(Conv)和来自特征图 和 的Sigmoid激活函数来捕获。这一信息用于校准特征图,并促进对显著空间区域的强调。
D-Net Architecture
D-Net的整体架构包括一个编码器、一个瓶颈层、一个解码器以及一个显著性层(图3)。显著性层用于从原始图像中提取显著的空间特征,而编码器-解码器架构负责学习层次化的特征表示。
3.2.2 Encoder.
作者不采用将块展平后使用线性层进行投影的方法,而是使用一个大型 的卷积,步长为2,将图像分割成尺寸为 的特征嵌入。这些特征嵌入随后被投影成 维向量()。在每一个阶段,作者结合两个连续的DLK块来提取上下文信息。
为了在降采样块中交换通道间的信息,作者使用一个 Kernel 大小为 且步长为2的卷积层来缩小特征图,并将通道数增加一倍。在每一个阶段,输出特征图的尺寸分别为 ,,,以及 。
3.2.3 Bottleneck.
两个连续的DLK块用于 Neck 分。输入和输出特征的空间维度均为 。
3.2.4 Decoder.
在每一个阶段,都使用一个步长为2的转置卷积来放大特征图,并通过2倍的因子减少通道数。这些上采样的特征随后通过DFF模块内的跳跃连接与编码器的特征进行融合。接着使用两个连续的DLK块。
在每个阶段的输出特征图的维度分别为,,,以及。最后,使用一个转置卷积层将特征图上采样到的维度。
2.3.2 Salience layer.
一个卷积块,它由两个连续的卷积层组成,用于从输入图像生成维度为的特征。这些特征在一个DFF模块中与解码器的特征进行融合。接着,另一个卷积块被用来捕捉更细致的特征。最后,使用一个卷积层来生成 Voxel 级的分割预测。
3 Experiments and results
3.0.1 Datasets.
作者进行了两项公开数据集的实验。第一项是MICCAI 2022 AMOS挑战赛数据集(AMOS 2022)。它包括300张多对比度腹部CT图像,其中有15个解剖器官为腹部多器官分割手动标注。通过MONAI4实现的处理流程将这些3D体数据预处理并增强为尺寸的 Voxel 块。
第二项是医疗分割十项全能(MSD)脑肿瘤挑战赛数据集。它由484个主体组成,每个主体有四种3D MRI模式(FLAIR、T1w、T1gd、T2w)和三种前景标注:水肿(ED)、增强型肿瘤(ET)和非增强型肿瘤(NET)。数据通过nnUNet处理流程[11]预处理为尺寸的 Voxel 块。
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。
表1展示了在AMOS腹部多器官分割任务上的性能对比。D-Net以相对较少的FLOPs和最低的参数数量取得了最佳的整体性能。此外,D-Net在所有特定器官分割任务中的Dice得分都有显著提升。
表2展示了MSD脑肿瘤分割任务的结果。与其它分割方法相比,D-Net在所有分割任务上都展示了优越的性能。
3.3.2 Ablation study.
为了进行消融研究,作者将D-Net解构为DLK-Net,方法是移除D-Net中的显著层,并将每个DFF模块替换为连接层后面跟着一个卷积层。与其他 Baseline 方法相比,DLK-Net在两个分割任务中均显示出更高的分割精度,同时在模型复杂度上最低(表1和表2)。
4 Conclusion
作者引入了D-Net用于体积分割医学图像,通过将动态大核模块和动态特征融合模块整合到分层 Transformer 架构中。动态大核块被采纳为基本块,用于通用多尺度局部特征提取和自适应的全局空间信息利用。
此外,还提出了动态特征融合模块以实现自适应特征融合。在两个分割任务上,即腹部多器官分割和脑肿瘤分割,D-Net的表现优于当前流行的 Baseline 。作者相信D-Net有潜力在各类医学图像分割任务上实现令人期待的分割性能。
参考
[1].D-Net: Dynamic Large Kernel with Dynamic Feature Fusion for Volumetric Medical Image Segmentation.
点击上方卡片,关注 「AI视界引擎」 公众号