点击下方卡片,关注 「AI视界引擎」 公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
随着深度学习的快速发展,基于卷积神经网络(CNN)的U形网络在医学图像分割方面取得了成功,并广泛应用于各种任务。然而,它们在捕捉全局特征方面的局限性,限制了它们在复杂分割任务上的性能。
视觉 Transformer (ViT)的出现有效地弥补了CNN的这一缺陷,并推动了基于ViT的U形网络在医学图像分割领域的应用。
然而,ViT的高计算需求使其不适合许多具有有限资源的医疗设备和移动平台,限制了其在资源受限的边缘设备上的部署。
为了解决这个问题,作者提出了Evit-UNet,这是一种高效的基于ViT的分割网络,它降低了计算复杂性,同时保持了准确性,使其成为资源受限的医疗设备的理想选择。
Evit-UNet是基于U形结构的,包括编码器、解码器、 Bottleneck 层和 Shortcut ,结合卷积操作和自注意力机制以优化效率。
实验结果表明,Evit-UNet在医学图像分割方面实现了高精度,同时显著降低了计算复杂性。
1 Introduction
随着深度学习的快速发展,医学图像分析领域取得了显著进步,尤其是在图像分割任务方面。U-shape网络已成为分割网络的主流模型。这种结构将对称编码-解码配置与 Bottleneck 层和 Shortcut 相结合,构成了经典的U-shaped分割网络[1]。经典的U-Net架构采用卷积下采样来捕捉多级特征。它通过 Shortcut 来保留空间细节,然后通过基于解码器的上采样来恢复图像分辨率以实现精确像素级分割。
这种结构在各种医学图像分割任务上取得了很大的成功,如心脏、器官和病变分割。一些基于U-Net的工作旨在提高网络性能,例如U-Net++,它通过优化 Shortcut 机制来提高网络效率。此外,DeepLabV3[3]通过优化卷积操作来增强多尺度特征的处理。然而,基于CNN的UNet仍需在捕捉全局语义信息和处理复杂特征交互方面进行改进。
视觉 Transformer (ViTs)的引入解决了卷积神经网络(CNNs)在捕捉全局信息方面的局限性。ViTs 中的自注意力机制有效地捕获了全局上下文 [4],在处理全局特征和长程依赖关系方面相较于传统 CNNs 具有显著优势。研究行人已经开始探索其在医学图像分割领域的应用。
例如,Att-UNet [5] 和 TransUNet [6] 将 Transformer 引入到 UNet 网络中,而混合 CNN-Transformer 模型如 HiFormer [7] 和 UCTransUNet [8] 则被设计出来以提高网络速度。
此外,SwinUNet [9] 和 MedT [10] 利用专门的Transformer架构来提高网络速度和性能。ViTs中的自注意力机制提高了分割的准确性和鲁棒性,进一步推动了在医学图像分割任务中的效率。
尽管ViTs在视觉任务上表现出色,但它们的高计算复杂性限制了在资源受限设备上的应用 [11, 12]。在医学图像分割任务中,在保持准确性的同时降低计算复杂性一直是资源受限设备分割任务的关键追求 [9, 7]。
因此,作者提出了EViT-UNet,这是一种基于ViT的高效U形网络,用于在移动和边缘设备上的医学图像分割。它通过结合卷积和自注意力机制,在降低计算复杂性的同时继承了ViT捕获全局信息的能力,确保了高精度,同时最小化了计算成本,使其成为移动和边缘设备的理想选择。在多个数据集上进行测试,EViT-UNet展示了卓越的分割精度,并超过了其他流行的分割框架。
作者的关键贡献如下:
(1) 开发了一种基于ViT的高效U形分割框架,该框架集成了一个编码器、一个解码器以及 Shortcut ,并在多种数据集上表现出卓越的性能。
(2) 在与多个网络的比较分析中实现了最佳的计算效率。
(3) 在保持高准确性的同时成功降低了计算复杂性,从而提高了在资源受限环境中部署这种技术的可行性,用于医学图像分割任务。
2 Method
Architecture overview
作者的网络总体架构如图2(a)所示。设计包括编码器、解码器、 Bottleneck 层和上采样过程中的 Short-Cut 。编码器和解码器均分为四个阶段,作者采用EfficientFormerV2块[16]作为基本单元。通过初始特征提取块(stem)和下采样到尺寸tp 。
然后,在每个编码器阶段的块下采样,下采样率为2。编码器采用图2(a)所示的通道配置。编码器的最后一阶段进行全局特征融合,并将特征传递给解码器。作者设计了一个与编码器对称的解码器。解码器特征通过 Short-Cut 与编码器特征结合,恢复图像特征,并在每个阶段进行2x上采样。最后,上采样模块进行4x上采样并输出像素级预测。
Efficientformer block
与仅使用Transformer和自注意力作为编码器的模型不同,作者的网络采用了一种混合方法,结合卷积和自注意力模块。在高分辨率阶段,自注意力需要计算所有像素之间的交互,导致显著的计算开销[4]。
因此,作者的模块使用深度卷积(DW)卷积[13]构建 FFN (FFN)以提取局部特征,如图2(b)所示。与标准卷积相比,DW卷积对每个输入通道应用一个滤波器,显著降低了计算复杂性并增强了局部特征。这个过程可以描述为:
其中, 是第 层,在 阶段,而 是一个可学习的层缩放 [17]。
在低分辨率阶段,自注意力机制的计算负担显著降低。作者的模块引入了多头自注意力(MHSA)机制4,增强了捕捉全局特征的能力,丰富了编码器中的多尺度特征。在解码器中,多头注意力机制通过聚合全局和局部特征提高了图像重建的准确性。这种方法有效地平衡了准确性和计算效率,使模型能够捕捉复杂的全局依赖性,而不会显著增加计算负载。这个过程可以描述如下:
输入特征通过映射函数 Proj() 经过线性变换后,得到注意力机制中的 Query (Q)、键(K)和值(V):
位置编码的 learnable 注意力偏差为 。
Downsampling and Upsampling
在降采样过程中,卷积在高低分辨率阶段中都被采用,以实现高效的降采样和减小特征图的大小。在低分辨率阶段,作者采用自注意力机制进行降采样,通过调整 Query Token 的数量,可以有效地捕捉降采样过程中的全局依赖和多尺度特征。
这种方法平衡了降采样过程中自注意力机制的计算复杂度,因为分辨率降低。作者在解码器中提出了一个对称设计,通过调整 Query Token 的数量,在低分辨率阶段利用自注意力进行上采样。同时,卷积操作被用于上采样高分辨率阶段。虽然保证图像重建的准确性,但同时降低了计算复杂度。自注意力降采样/上采样可以描述为:
当时,它表示下采样;当=2时,它表示上采样。
Skip Connection
在U形分割网络中,跳接连接(skip connections)起着关键作用,通过将编码器收集到的特征传递给解码器,有效地保留低级特征[1]。然而,最近的研究发现传统跳接连接(skip connections)存在一些局限性[8]。简单地将编码器和解码器特征拼接在一起可能会引入冗余,而且由于跳接连接主要传递局部特征,它们在更复杂的分割任务中难以捕捉全局依赖性。
然而,一些研究将注意力机制引入跳接连接,并取得了令人瞩目的成果[8, 5]。基于调查和研究,作者在跳接连接中引入了通道注意力(channel attention)[8]。这种方法通过应用注意力机制来强调重要的特征通道并抑制冗余,同时有助于在不同特征尺度上建立更好的全局依赖建模,而无需增加额外的计算开销。基于通道的跳接连接可以描述如下:
这里, 是前一层特征图, 是来自编码器在降采样过程中的 Shortcut 特征图。
3 Experiments and Results
Synapse多器官分割数据集(Synapse)[19]包含30个案例,总计3779个轴向腹部临床CT图像。该数据集分为18个训练样本和12个测试样本。作者的方法在包括主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺和胃在内的八个腹部器官上进行评估。
GlaS数据集(GlaS)[20]和MoNuSeg数据集(MoNuSeg)[21] 分别包含85张训练图像、80张测试图像,以及30张训练图像和14张测试图像。作者在GlaS和MoNuSeg数据集上进行5折交叉验证。
实验实现作者的网络基于Python 3.10和PyTorch 2.0实现。输入图像大小为224x224,批量大小为32,并在单个Nvidia A100 GPU上进行训练。作者采用了来自ImageNet的预训练权重EfficientFormerV2,并将它们适应并加载到编码器和解码器中,以尽可能地初始化模型参数。使用SGD优化器进行反向传播优化。
表1显示了在Synapse数据集上的结果,表2显示了在Glas和MoNuSeg数据集上的结果。作者的模型在许多当前方法中脱颖而出,在80.87%的性能上实现了优越的性能。具体而言,它在Synapse数据集上的最佳DSC(Deep Sentiment Classification)平均提高了0.33%。Glas和MoNuSeg数据集上的结果表明,作者的方法在这两个数据集上表现良好。
具体而言,在Glas数据集上,作者的模型实现了最佳的DSC(Deep Sentiment Classification)92.44%和IOU(Intersection over Union)86.50%。在MoNuSeg数据集上,作者的模型也实现了DSC(Deep Sentiment Classification)79.27%和IOU(Intersection over Union)65.87%,超越了许多流行的比较方法。
作者还获得了Synapse数据集(图3(A))、Glas数据集(图3(B))和MoNuSeg数据集(图3(C))的视觉结果,以说明作者方法的表现。最重要的是,作者在Synapse数据集上进行了计算复杂性的比较;作者的方法在计算效率方面超越了所有比较方法,计算复杂性仅为6.39 GMac(图1)。
4 Conclusion and Discussion
综上所述,作者开发了一种分割框架,该框架在性能方面表现出色,同时具有优越的计算效率。作者的模型在准确性方面超过了众多最先进的方法,同时保持了较低的计算负担,使其特别适合于计算资源有限的设备,如医疗设备。这些特性强调了该模型在性能关键的真实世界应用中的适用性。
尽管该模型表现出卓越的性能和高效率,但其在医疗设备方面的适应性和实用性仍存在一些局限。医疗设备有复杂的要求,虽然作者的框架在实验中表现良好,但为了在便携式和嵌入式系统中更广泛地部署,仍需进一步优化。展望未来,作者的研究可以专注于针对特定硬件实现的模型微调,从而提高其在实际医疗设备中的适用性。
5 Compliance with Ethical Standards
该研究采用公开获取的人体数据,通过Synapse [19]、GlaS [20] 和 MoNuSeg [21] 进行回顾性研究。由于公开获取数据无需伦理批准,因此无需获得伦理批准。
[0]. Evit-Unet: U-Net like efficient vision transformer for medical image segmentation on mobile and edge devices.
点击上方卡片,关注 「AI视界引擎」 公众号
