备注好友:方向-学校/公司-姓名/昵称
【AIGC 先锋科技】交流群
U-Net及其变体在医学图像分割中应用广泛。然而,这些模型的大量参数和计算复杂性使得它们在计算资源有限的临床环境中使用不太合适。
为了解决这一限制,作者提出了一种新型的轻量级位移U-Net(LSU-Net)。作者以轻量级的方式整合了轻量级卷积块和 Token 化位移块,并将它们与动态权重多损失设计相结合,以实现高效的动态权重分配。轻量级卷积块通过结合标准卷积和深度可分离卷积,以较低参数数量有效捕捉特征。
Token 化位移块通过空间位移块和深度可分离卷积的组合,优化了特征表示,并通过位移和捕捉深度特征。在每个层动态调整损失权重,接近最优解,并增强训练稳定性。作者在UWMGI和MSD结肠数据集上验证了LSU-Net,实验结果表明,LSU-Net在大多数最先进的分割架构中表现更优。
医学图像分割对于提取结构细节以辅助诊断和治疗至关重要。深度学习通过多样化的网络架构显著推动了这一领域的发展,尤其是从UNet [1] 开始的U型编码器-解码器设计。随后,如
[2]、Attention U-Net [3]、UNet3
[4]、Rolling-UNet [5] 和 UNeXt [6] 等模型,引入了嵌套结构和密集跳接等创新,进一步提升了性能。
受到Transformer在自然语言处理和计算机视觉领域成功应用的启发,基于Transformer的网络,如Swin-UNet[7]和TransUNet[8],能够捕捉长距离依赖关系进行分割。MedT[9]和TransBTS[10]等模型通过先进的注意力机制[11, 12]提高了准确率。
尽管取得了这些进步,许多模型却忽略了实际医疗场景中的计算限制。先前的研究也难以在这些限制内平衡参数和性能。为了解决这些挑战,作者提出了LSU-Net,这是一种轻量级的分割模型,在执行多尺度监督学习的同时保持强大的分割性能。LSU-Net保留了经典的编码器-解码器架构,在不同阶段应用不同的模块。在浅层阶段,它采用低参数卷积模块来高效地捕捉特征。在深层阶段,LSU-Net使用MLP(多层感知器)来转换深层特征,捕捉区域依赖关系并优化特征表示。作者还实现了横向深度监督,动态调整损失权重以比较多尺度信息。
为了验证LSU-Net的有效性,作者在UWMGI [13]和MSD结肠[14]数据集上对其进行了评估。结果显示,作者的方法在参数较少的情况下实现了出色的分割性能,证明了其在实际环境中的应用性。
- 方法
本研究采用了一系列先进的机器学习算法来分析和处理数据。具体方法如下:
- 数据收集:作者从多个来源收集了大量的数据集,包括公开数据库和内部数据源。
- 数据预处理:对收集到的数据进行清洗和标准化处理,以确保数据质量。
- 特征工程:通过提取和构造特征,将原始数据转换为适合机器学习模型输入的形式。
- 模型选择:根据研究目标,作者选择了多种机器学习模型,包括支持向量机(SVM)、随机森林(RF)和深度神经网络(DNN)。
- 模型训练:使用预处理后的数据对选定的模型进行训练,调整模型参数以优化性能。
- 模型评估:通过交叉验证等方法对模型的性能进行评估,并选择最佳模型。
- 结果分析:对模型输出结果进行深入分析,以揭示数据中的潜在模式和规律。
- 模型优化:根据分析结果对模型进行进一步优化,以提高预测准确性和泛化能力。
本节详细介绍了LSU-Net的相关信息,包括其整体网络架构、核心组件和损失函数。讨论的核心组件包括轻量级卷积块和 Token 化平移块。
整体架构。作者提出了一种名为LSU-Net的轻量级网络,该网络旨在通过高效的特徵提取和处理进行腹部器官分割。LSU-Net采用U型架构,如图1所示。为实现这一目标,作者引入了两个模块:轻量级卷积块和 Token 化位移块。编码器包括编码器轻量级卷积块和 Token 化位移下采样块,而解码器则由解码器轻量级卷积块和 Token 化位移上采样块组成。
轻量级卷积块。如图1左下角所示,轻量级卷积块由多层组成:首先,主分支通过3×3卷积来提取局部特征并利用3×3核捕捉空间信息。接下来,批量归一化(BN)和GELU激活函数增强了特征表示,如公式(1)所示:
在此,
代表输入,
表示该阶段的输出。
在此之后,使用深度可分离卷积(DwSConv)进行特征提取,显著降低了计算量和参数数量。接着,应用分组归一化(GNorm)。随后,一个
的卷积调整特征图的通道数,GELU激活增强了非线性表示。最后,编码器中对特征图的空间维度进行下采样的是步长为2的
卷积,而在解码器中应用的是步长为1的
卷积。
此外,卷积块包含多种 Shortcut 。编码器使用
卷积作为 Shortcut 来改变通道维度。编码器和解码器的输出分别由公式(4)和公式(5)给出。这种设计集成了多种卷积和归一化操作,以平衡特征提取、计算效率和模型性能。
分词位移模块。分词位移模块沿着通道维度扩展特征图,有效地捕捉空间和通道特征,从而提升模型的图像分割性能。该模块的结构如图1左上角所示。该结构首先使用一个
卷积层处理输入,以调整特征通道。同时,初始输入特征图被送入空间位移块[15],该块通过分词机制进一步处理输入特征。空间位移块沿着通道维度扩展特征图,然后将扩展后的特征图分为三部分。其中一部分保持静止,而另外两部分则在两个方向上进行空间位移操作,如图2所示。图2(a)对应于公式(6):
在此,
代表单个样本,而
、
和
分别代表高度、宽度和通道。图2(b)对应于公式(7)。
不同的空间平移操作被应用于这些分割的部分,随后通过分区注意力机制进行融合。随后,原始特征图和经过处理的特征图被融合以增强特征表示,如图(8)所示。
在本文中,“shift”指的是空间位移块(Spatial Shift Block)。随后,使用了一个深度可分离卷积层(DWsConv layer),在保持卷积操作有效性的同时,减少了参数和计算复杂度。接着,通过一个全局归一化层(GNorm layer)对特征图进行归一化处理,从而加速训练过程并提升模型的泛化能力。
随后,应用一个具有GELU激活函数的
卷积层,以进一步执行非线性变换和特征提取。最后,另一个
卷积层对特征进行处理和融合,输出最终的特征图。该过程由公式(10)计算得出。
Tokenized Shift Block的下采样模块在输入端增加了一个重叠 Patch 嵌入层。该层将输入图像划分为小 Patch 并进行嵌入表示,从而保留了更多局部信息,增强了特征表示的丰富性。
损失函数。作者提出了多尺度深度损失(MDL),该损失函数平衡了类别特征并减少了模糊边缘的影响,以提升精度。MDL在不同的尺度下对 Mask 进行下采样,为每个U-Net Level 分配相应的 Mask 。损失是通过分割结果计算的,每个 Level 使用不同的损失函数。最后,自动加权损失(AWL)[16]函数结合了多级损失,根据公式(11)更新权重。
在此,
表示损失函数,
是多尺度输出的总层数,
表示不同尺度的损失,而
则代表不同层级的损失函数的自适应权重。项
限制了权重参数的过度膨胀,防止在优化过程中出现极端值。这种方法增强了损失函数的稳定性,并允许模型根据训练数据自动调整不同损失的权重。
- 实验
在本节中,作者对所提出的网络架构在UWMGI[13]和MSD Colon[14]数据集上进行验证。首先,作者概述了一些实现细节。然后,作者将LSU-Net与其他最先进的架构在这两个数据集上进行比较,并分析实验结果。
实施细节。作者使用Pytorch [17]框架开发LSU-Net。所有实验均在单个NVIDIA GeForce RTX 3090 GPU上运行。为了评估作者的方法,作者采用平均交并比(mIoU)和Dice相似系数(DsC)作为指标。mIoU和DsC的计算公式如下:
表示类别数量,
和
分别表示预测值和标签。
图像被归一化并调整至
像素大小用于UWMGI数据集,以及
像素大小用于MSD结直肠癌数据集。采用Adam[19]作为优化器,初始学习率为0.001,并使用CosineAnnealingLR[20]作为调度器。调度器配置了最大迭代次数和周期,以及最小学习率为1e-5。两个数据集都训练了100个周期,其中UWMGI数据集的批次大小为16,MSD结直肠癌数据集的批次大小为10。作者使用平均交并比(mIoU)和Dice相似系数(DsC)来评估作者的方法。
实验结果。表1中的结果显示,LSUNet在UWMGI数据集上优于其他模型。LSUNet超越了 Baseline 模型UNeXt,实现了3.75%的mloU和3.41%的DSC提升。与Res-UNet相比,LSU-Net不仅获得了6.98%的mIoU和6.73%的DSC提升,而且分别减少了12.1倍和47.4倍的参数和计算量。LSU-Net在mloU和DSC方面也优于UNet++,分别提高了2.71%和2.78%,同时减少了32.3倍和46.4倍的参数和计算量。在MSD结肠癌数据集上,如表2所示,LSUNet在性能和参数之间取得了平衡,提供了出色的分割结果。图3展示了具体的分割示例。
表3. 消融研究检验了每个模块在UWMGI数据集上的贡献。
表4. 消融研究检验了每个模块对MSD结直肠癌数据集的贡献。
作者对UWMGI和MSD结直肠癌数据集进行了消融研究,以评估每个模块对监督学习的影响。表3显示,引入MDL显著提升了性能,与未使用MDL的模型相比,UWMGI上的mloU提升了1.23%,DSC提升了1.18%。移除其他模块同样会降低性能。表4显示,这种效果在MSD结直肠癌数据集中也得到了观察。结果表明,每个模块的贡献各不相同,与单类的MSD结直肠癌相比,MDL对多类的UWMGI数据集的影响更为显著,这表明需要针对数据集进行特定的微调。Light Conv Block和Tokenized Shift Block的消融表明,这两个模块在不同的数据集中均能提升性能。
本文提出了一种名为LSU-Net的轻量级腹部器官分割网络,该网络利用特征空间位移和多尺度损失来增强器官区域的识别,同时采用轻量级卷积来减少模型参数。
作者在多个腹部数据集上进行了广泛的实验,结果表明所提出的LSUNet优于现有的分割模型,并推动了前沿性能的提升。作者的模型在临床应用中具有巨大潜力,例如量化病变面积和监测疾病进展。
参考
[1]. LSU-NET: LIGHTWEIGHT AUTOMATIC ORGANS SEGMENTATION NETWORK FORMEDICAL IMAGES .