一种轻量级、内存效率高的3D U-Net模型,结合注意力机制和不确定性建模技术,用于冠状动脉分割
《Computer Methods and Programs in Biomedicine Update》:A lightweight memory-efficient 3D U-Net with attention and uncertainty modeling for coronary artery segmentation
【字体:
大
中
小
】
时间:2026年03月25日
来源:Computer Methods and Programs in Biomedicine Update CS9.8
编辑推荐:
郝迪·权 | 帕诺斯·利亚蒂西斯
计算机科学系,哈利法大学,邮政信箱127788,阿布扎比,阿拉伯联合酋长国
**摘要**
深度学习(DL)方法已被广泛应用于3D冠状动脉分割。然而,由于3D CT图像的数据量庞大,高效的GPU内存利用已成为当前DL分割模型实际应用的
郝迪·权 | 帕诺斯·利亚蒂西斯
计算机科学系,哈利法大学,邮政信箱127788,阿布扎比,阿拉伯联合酋长国
**摘要**
深度学习(DL)方法已被广泛应用于3D冠状动脉分割。然而,由于3D CT图像的数据量庞大,高效的GPU内存利用已成为当前DL分割模型实际应用的主要瓶颈。此外,在特征图中有效捕捉管状结构信息同时抑制背景噪声仍然是一个重大挑战。为了解决这些问题,我们提出了一种内存高效的U-Net(ME-UNet)——这是一种轻量级的U-Net变体,它去除了每个块中的残差连接,并引入了两个关键组件:不确定性跳跃连接(USC)和维度融合注意力(DFA)。实验结果表明,ME-UNet不仅降低了GPU内存消耗,还实现了最先进的分割性能。源代码可在以下链接获取:https://github.com/hd1437/ME-UNet
**1. 引言**
冠状动脉疾病(CAD)是最常见的心血管疾病,具有高发病率和高死亡率,每年导致全球超过四分之一的死亡病例[1]。冠状动脉CT造影(CCTA)是评估冠状动脉解剖结构和检测斑块的首选无创成像方式,因为它在检测CAD方面具有高特异性和敏感性[2]。冠状动脉树的复杂解剖结构和冠状动脉的精细形态使得手动处理非常耗时,并需要高级技术专长[3]。随着对CCTA程序需求的稳步增长,医疗保健提供者在有效满足患者需求方面面临着越来越大的挑战。这给成本上升和现有资源的利用带来了压力,影响了服务效率和患者的及时诊断与护理。因此,迫切需要开发准确且自动化的冠状动脉分割方法来支持放射科医生和临床医生。近年来,已有大量研究使用传统的图像处理技术来支持血管分割,例如图像滤波和增强算法[4]、区域生长方法[5]等。然而,大多数基于图像处理和机器学习的经典冠状动脉分割方法依赖于手工设计的特征工程,因此在处理丰富的血管形态变化方面存在困难[6]。最近,深度学习(DL)范式应运而生,它具有自动学习复杂数据模式的强大能力,在文本、语音、图像和视频处理等各种应用中变得无处不在。这导致了深度学习模型在医学成像中的广泛应用,包括CT冠状血管分割模型。与传统血管分割方法相比,DL方法能够自动捕捉复杂的图像特征,充分利用了它们的表示能力。例如,基于卷积神经网络(CNN)架构的U-Net[7]在血管分割方面取得了良好的效果。然而,尽管这些方法在2D血管分割任务中表现出色,但在3D血管分割任务中需要更多的计算资源,并且推理时间较长。
如2.1节所讨论的,3D血管分割一直是许多DL研究的重点,从而显著提高了性能。然而,在实际临床应用中,仅靠性能是不够的。深度学习模型通常部署在单GPU工作站或共享计算服务器上,这些设备的GPU内存有限,而不是在高端研究级硬件上[8]。在这种资源受限的环境下,许多最先进的3D分割模型由于过高的GPU内存需求而难以集成到常规临床工作流程中。在最先进的DL系统中,GPU内存的使用主要由编码器/解码器块中的残差连接和注意力机制的复杂性驱动。虽然这些组件对模型性能很重要,但跳跃连接也起着关键作用——尽管其内存开销要低得多。在这项工作中,我们有意减少对残差连接和复杂注意力的依赖,以提高内存效率。因此,如2.2节进一步讨论的,增强跳跃连接的有效性对于维持分割性能至关重要。基于上述启发,本研究旨在开发一种准确且高效的3D血管分割模型,即内存高效的UNet(ME-UNet),以适应资源受限的医疗环境中的临床工作流程。具体贡献如下:
1. 提出了一种轻量级的U-Net骨干网络,用于3D冠状血管分割,旨在实现分割精度和GPU内存消耗之间的良好平衡。这是通过简化编码器-解码器结构并减少对残差连接的依赖来实现的,从而在保持性能的同时显著降低计算和内存成本。
2. 引入了两种高效的特征增强模块:(i)不确定性跳跃连接(USC),它对体素级别的不确定性进行建模,以在跳跃路径中抑制噪声背景特征,同时具有适度的记忆开销;(ii)维度融合注意力(DFA),这是一种轻量级的注意力机制,能够以线性复杂度捕获互补的通道级和空间信息,进一步提高分割精度,而不会显著增加GPU内存使用量。
3. 实验结果表明,所提出的框架在运行时间和GPU资源消耗方面具有优势,同时提供了高性能。与最先进的UNet++相比,ME-UNet将AUC提高了0.011,DSC提高了0.007,IoU提高了0.015。此外,UNet++的GPU内存消耗为18.4 GB,而ME-UNet仅需4.7 GB。
**2. 相关工作**
2.1. 3D冠状动脉分割
如前所述,CCTA中准确的3D冠状动脉分割对于冠状动脉疾病的诊断至关重要。在本小节中,我们将概述最先进的方法,从扩展U-Net架构到处理3D体积,再到整合注意力机制。
3D U-Net[9]是一种经典的血管分割模型,采用基于CNN的编码器-解码器结构。它在捕获局部结构信息方面很有效;然而,在建模图像的全局结构上下文方面存在局限性[10]。为了克服这一限制,Pan等人[11]提出了一种模型,将U-Net和Transformer模块之间的多尺度特征交互结合起来,实现了局部和全局结构信息的整合。Wu等人[12]进一步在U-Net中引入了自注意力块,使模型在处理每个体素时能够考虑整个图像的信息,从而增强了其对全局关系的理解。然而,3D医学图像的大规模对基于注意力的模型的效率提出了挑战。传统注意力机制的局限性在于需要计算每个体素与图像中所有其他体素之间的注意力得分,这导致了过高的计算和内存成本[13]。为了解决这个问题,Perera等人[14]提出了SegFormer3D,它对多尺度体积特征进行注意力计算,有效平衡了局部细节保留和全局语义建模,同时显著降低了计算开销,相比传统的大规模3D Transformer。同样,Ren等人[15]引入了HResFormer,它结合了2D和3D Transformer分支,并采用混合局部-全局融合模块,有效地结合了细粒度的切片内特征和切片间上下文信息,从而提高了对体积结构的理解。尽管这些方法在分割性能上取得了显著改进,但它们大多仍然依赖于计算密集型的注意力机制,导致大量的GPU资源消耗,限制了其在临床应用中的潜力。
2.2. 跳跃连接
U-Net中的跳跃连接是影响图像分割质量的关键因素。它们通过添加额外的连接来链接输入和输出信息,有效缓解了网络层之间信息传播时梯度消失或爆炸的问题[16]。然而,传统的跳跃连接通常直接压缩局部信息,引入了过多的无关背景噪声[17]。为了解决这个问题,Dong等人[3]用一种新的基于注意力的特征融合模块替换了传统的跳跃连接。Wu等人[18]引入了双注意力块来捕获长距离特征。然而,这些基于注意力的跳跃连接方法计算成本较高,因为它们必须计算每对体素之间的相关性,并且在反向传播过程中需要存储额外的梯度信息来计算前一层的梯度[19]。大多数先前的研究主要集中在提高网络性能上,这导致了基于注意力的跳跃连接的过度使用,从而导致了大量的GPU内存消耗。因此,我们提出了一种轻量级的卷积跳跃连接USC,以有效保留关键特征信息,同时显著减少GPU内存使用量。
2.3. 注意力机制
3D血管分割的注意力机制研究可以分为三类:通道注意力、空间注意力以及通道和空间注意力相结合的注意力。通道注意力允许模型关注输入数据的不同通道,从而更好地捕获每个通道的关键信息[20]。Wang等人[21]将通道注意力机制应用于跳跃连接,以传递信息并选择有用的特征图。对于空间注意力,Guo等人[22]引入了一种空间注意力模块,沿空间维度推断注意力图,并将其与输入特征图相乘以实现自适应特征细化。一些研究结合了这两种类型的注意力机制,以充分利用它们的优势。Mou等人[23]提出了通道和空间注意力模块,对通道和空间特征图都应用了自注意力。Xiao等人[24]提出了RDTransUNet,这是一种混合U-Net架构,它结合了残差连接和双通道-空间注意力机制,旨在共同捕获局部细节和全局上下文信息,用于医学图像分割。尽管之前提出的注意力机制在捕获全局结构信息方面很有效,但它们通常会产生较大的计算开销[25]。为了解决这一限制,我们提出了一种新的轻量级注意力机制DFA,以在尽可能保留全局信息的同时减少GPU内存使用量。
**3. 方法论**
原始的3D血管图像表示为X0,相应的标签体积表示为Y。编码器在四个层次生成的特征图分别表示为X1、X2、X3、X4。解码器生成的相应重建图像分别表示为D0、D1、D2、D3、D4。这里,X0、D0∈R1×H×W×D,Xi、Di∈R23+i×H2i×W2i×D2i。
**3.1. 框架**
经典的3D U-Net使用卷积块处理空间信息。虽然在处理局部空间特征方面很有效,但在捕获全局结构信息方面存在局限性。为了解决这一限制,一些研究用基于注意力的块替换了卷积块[26]、[27]、[28]、[29]、[30]。然而,这些方法通常在整个网络中引入全局自注意力或复杂的令牌交互,这不可避免地增加了架构复杂性和GPU内存消耗,特别是对于3D体积。为了提高GPU内存效率,本研究采用了轻量级的基于注意力的U-Net骨干网络。所提出的ME-UNet在最深的编码器-解码器层中结合了维度融合注意力(DFA),实现了高效处理,同时保留了全局信息。所提出架构的概览如图1所示。
**图1. 所提出的ME-UNet框架的整体架构。**该网络采用编码器-解码器设计,在最深的编码器-解码器层应用了维度融合注意力(DFA)模块,以低计算成本高效捕获全局结构信息。此外,在编码器和解码器之间引入了不确定性跳跃连接(USC),以抑制特征图中的噪声信息,提高血管分割的鲁棒性和泛化能力。
**3.2. 维度融合注意力**
为了有效建模全局结构和通道级依赖性,提出了维度融合注意力(DFA)模块。与现有的在空间维度上应用均匀或对称注意力的通道-空间注意力公式不同,DFA沿三个正交轴分解空间注意力,并将其与轻量级的通道交互相结合,明确针对3D体积中的内存效率。对于通道注意力,使用两个大小为C×C的可学习矩阵在特征级别编码通道间关系,避免了体素级的自注意力,从而防止了内存消耗的二次增长。通道注意力的操作在公式(1)中表述:
(1) Ca(Xi) = R?1(R(Xi) Softmax(R(Xi)?(R(Xi)W1)W2))
这里,运算符R(?)将特征张量Xi∈RC×H×W×D重塑为2D矩阵R(Xi)∈RN×C,其中N=H?W?D,C表示通道数量。W1、W2∈RC×C是可学习矩阵。因此,R(Xi)?(R(Xi)W1)产生一个大小为RC×C的通道相关矩阵。经过Softmax操作后,将得到的注意力矩阵应用于R(Xi)以重新加权通道响应,输出通过R?1(?)重塑。
对于空间注意力,沿3D体积的三个正交方向应用卷积滤波器生成方向查询向量。这些查询向量与来自卷积操作的相应关键向量逐元素相乘。得到的特征图沿着通道维度进行拼接,形成一个组合的查询-键表示,然后与卷积层产生的值向量进行逐元素相乘。由于逐元素乘法的空间复杂度为O(n),所提出的DFA模块在内存使用上非常高效。详细公式在方程(2)中给出。(2)Qx=R(Conv1×3×3(X4)), Qy=R(Conv3×1×3(X4)), Qz=R(Conv3×3×1(X4)), K=R(Conv3(X4)), V=R(Conv1(X4)), Sa(X4)=R?1(Softmax(Qx?K⊕Qy?K⊕Qz?K)V)。这里,Qx、Qy和Qz分别是沿x、y和z空间维度的查询向量,K是键向量,Sa是空间注意力机制的最终输出,其中注意力分数用于加权值向量V。结合通道注意力和空间注意力使得模型能够全面理解血管结构的体积信息。通道注意力、空间注意力和原始输入的输出沿着通道维度进行拼接,然后输入到一个卷积层中,如方程(3)所示。(3)D4=Conv3×3×3(Ca(X4)⊕Sa(X4)⊕X4)。下载:下载高分辨率图像(502KB)下载:下载全尺寸图像图2. 所提出的DFA模块的架构。3.3. 不确定性跳跃连接跳跃连接直接将高分辨率的编码器特征传递给解码器,在这种情况下,特征融合特别容易受到噪声放大和血管边界不明确的影响[31]、[32]。在血管分割中,这种不可靠的低级特征在解码过程中可能会被放大,从而显著降低重建质量。因此,直接连接编码器和解码器特征容易导致过拟合,而基于注意力的跳跃连接通常会带来较大的计算和内存开销。为了解决这些问题,我们引入了不确定性跳跃连接(USC),它将不确定性估计直接集成到跳跃路径中。与在输出级别建模不确定性或依赖于变分推断的贝叶斯或概率分割方法[33]、[34]、[35]不同,USC仅在跳跃连接内部执行不确定性感知的调制,使得不可靠的编码器特征在合并到解码器表示之前可以被自适应地减弱。这种设计允许不确定性指导特征传播,而不是事后进行置信度估计,并且不引入推理时的采样。为了避免损坏细小的分支血管结构,首先应用一个卷积块来生成中间特征表示Hi:(4)Hi=ConvBlock(Xi)。然后,使用两个卷积层来估计3D特征体积中每个体素的高斯分布的均值μ和标准差σ。μ和σ在训练过程中端到端学习,以分别建模预期的特征响应及其相关的不确定性。μ和σ的每个元素对应于与单个体素相关的高斯分布的参数。均值μ使用3 × 3 × 3卷积来结合局部空间上下文,这对于保持细小的血管结构和捕获依赖于邻域的特征统计信息至关重要。相比之下,标准差σ代表特征响应的置信度而不是其空间结构,因此设计为较少依赖于局部邻域。因此,σ使用1 × 1 × 1卷积来估计,从而实现高效的通道级不确定性建模,而不引入额外的空间相关性。由于分布对标准差的大小敏感,应用Softplus函数来防止梯度爆炸:(5)μi=Conv3×3×3(Hi)(6)σi=Softplus(Conv1×1×1(Hi))。其中(7)Softplus(x)=ln(1+ex)。第二个隐藏表示Zi依赖于这两个参数。为了在反向传播过程中高效计算梯度,我们应用了重参数化技巧[36],它通过从参数化分布中采样而不是直接从输入空间采样,将随机元素?~N(0,1)引入模型。因此,Zi可以按照方程(8)计算:(8)Zi=μi+??σi。通过将Zi与重建图像Di+1拼接来生成重建图像Di。这样就得到了最终的Di输出。USC模块如图2所示。下载:下载高分辨率图像(151KB)下载:下载全尺寸图像图3. 所提出的USC模块的详细结构。3.4. DFAL的复杂性分析设N=C×H×W×D表示输入大小。通道注意力机制涉及重新塑造输入并执行矩阵乘法。重新塑造的输入大小为N。权重矩阵W1和W2的维度为C×C。因此,通道注意力的空间复杂度为O(N+C^2)。由于我们实验中C的最大值为128,远小于N,总空间复杂度简化为O(N)。对于空间注意力,卷积需要O(N?k^3),其中k是核大小。在我们的实验中,k设置为3。由于逐元素乘法和加法需要O(N),空间注意力的总空间复杂度也为O(N)。因此,DFA的整体空间复杂度为O(N)。表1总结了用于血管分割的几种注意力模块的空间复杂性,从理论上证明了所提出方法的效率。表1. 常用于体积医学图像分割的不同注意力机制的理论空间复杂性比较。注意力模块空间复杂性自注意力[37]O(N^2)分解自注意力[38]O(N^N)Reformer[39]O(Nlog(N))Linformer[40]O(N)DFAO(N)4. 实验4.1. 数据集在本研究中,我们使用ImageCAS数据集[41]进行训练,该数据集包含了2012年4月至2018年12月期间从广东省人民医院收集的1000个3D冠状动脉计算机断层扫描(CCTA)体积。所有1000个体积都用于我们的实验。每个体积的空间大小为512×512×(206–275)体素,平面分辨率范围为0.29至0.43 mm^2,层间距为0.25至0.45 mm。按照[41]中的预处理协议,我们对所有体积应用最近邻插值,统一将其大小调整为128 × 128 × 64,这显著减少了GPU内存消耗,同时保持了整体血管拓扑结构。数据集按患者分为训练集、验证集和测试集,分别包含700、100和200个体积。分割是固定的且对所有比较方法都是相同的,以确保公平比较,并且没有进行逐层分割以避免信息泄露。此外,我们还在公开可用的ASOCA数据集[42]上评估了所提出的方法,该数据集是作为MICCAI ASOCA挑战的一部分发布的。所有模型首先在ImageCAS数据集上进行训练,然后在ASOCA数据集上进行微调,以评估跨数据集的泛化性能。ASOCA数据集包含来自多个临床中心的40个CCTA体积,包括20个冠状动脉疾病病例和20个正常病例。原始体积的空间分辨率为512×512×(250–400)体素,不同扫描的体素间距各不相同。真实的冠状动脉注释由临床专家提供。为了确保与ImageCAS实验的一致性并实现公平的跨数据集评估,所有ASOCA体积都使用相同的重采样策略进行预处理,并统一大小为128 × 128 × 64。在40个病例中,25个用于微调,5个用于验证,10个用于测试。4.2. 实现细节我们使用AUC、Dice相似系数(DSC)和交并比(IoU)作为评估指标。这些指标通常介于0和1之间,其中较高的值表示更好的分割性能。AUC的计算公式为:(9)AUC=∫0^1TPR(t)d(FPR(t))。其中TPR、FPR和t分别代表真正例比率、假正例比率和相关阈值。DSC的计算公式为:(10)DSC=|Y?∩Y||Y?|+|Y|。IoU的计算公式为:(11)IoU=|Y?∩Y||Y?∪Y|。所有模型都是使用Python 3.9在安装了CUDA-12.1的RTX 4090 GPU上训练的。我们应用了Focal Loss[43]作为损失函数,其定义为:(12)FL(pt)=?αt(1?pt)γlog(pt),其中pt是真实类别的预测概率。平衡因子α和聚焦参数γ分别设置为0.25和5。在模型训练中使用了Adam优化器[44]。对于每个模型,我们重复实验30次以确保结果的可靠性和稳定性。我们将所提出方法的性能与七种最先进的方法进行了比较,即3D U-Net (2016) [9]、UNet++ (2018) [32]、SegResNet (2019) [45]、SwinUNETR (2021) [46]、UNETR (2022) [47]、CoT-UNet (2023) [12]和基于Transformer的3D U-Net (2023) [18]。3D U-Net、UNet++、SegResNet、SwinUNETR和UNETR模型是基于Monai [48]实现的,Monai是一个用于医学图像分析的开源Python库。CoT-UNet和基于Transformer的3D U-Net模型根据应用需求对其原始源代码进行了修改。为了确保实验之间的公平比较,我们使用Scikit-learn [49]中的网格搜索来确定所有模型的最佳学习率和最大训练周期。表2提供了训练过程的超参数详细信息。为了避免过度的内存消耗并确保更专注于内存使用情况,所有模型的批量大小都设置为1。表3显示了每个模型架构的具体超参数设置。表2. 每个分割模型的优化超参数设置。模型学习率最大周期3D-UNet0.00350UNet++0.00170SegResNet0.00280Swin-UNETR0.00260UNETR0.00170CoT-UNet0.00150基于Transformer的3D U-Net0.00370ME-UNet0.00160表3. 每个基线和所提出的ME-UNet模型的详细架构超参数设置。模型参数3D-UNet通道大小 = (16, 32, 64, 128, 256),残差单元数量 = 1UNet++特征 = (32, 64, 128, 128, 64),SegResNet向下块 = (8, 16, 32, 64),向上块 = (32, 16, 8)Swin-UNETR深度 = (2, 2, 2, 2),头数 = (3, 6, 12, 24),特征大小 = 120UNETR特征大小 = 64,隐藏层大小 = 128,多层感知器维度 = 512,头数 = 16CoT-UNet通道大小 = (16, 32, 64, 128, 256),上下文注意力维度 = 512基于Transformer的3D U-Net通道大小 = (16, 32, 64, 128, 256)双注意力维度 = ((16, 32), (32, 64), (64, 128), (128, 256))ME-UNet通道大小 = (16, 32, 64, 128)4.3. 实验结果表4和表5总结了所有比较模型的定量分割性能和计算效率。总体而言,ME-UNet在AUC、DSC和IoU方面在两个数据集上都始终表现出最佳的分割性能。在ImageCAS数据集上,ME-UNet获得了最高的AUC(0.847)、DSC(0.728)和IoU(0.583)。与最强的基线UNet++相比,ME-UNet将AUC提高了0.011(1.31%),DSC提高了0.007(0.97%),IoU提高了0.013(2.28%)。在ASOCA数据集上也观察到了类似的趋势,该模型首先在ImageCAS上训练,然后在ASOCA上进行微调。如表4所示,ME-UNet的表现优于所有基线,获得了AUC为0.862、DSC为0.754和IoU为0.605。这些结果表明,所提出的方法在跨数据集的微调设置下具有良好的泛化能力,突出了其在内存效率方面的稳健性。除了分割精度外,ME-UNet还在内存效率方面表现出优势,如表5所示。UNet++需要消耗18.4 GB的GPU内存,而ME-UNet将GPU内存使用量减少到仅4.7 GB(减少了74.5%)。这些结果突显了所提出架构所实现的显著效率提升,使其能够在标准的单GPU临床工作站上部署。表4. ImageCAS和ASOCA数据集上的定量分割性能。方法ImageCASASOCA空单元AUCDSCIoUAUCDSCIoU3D U-Net0.8020.6910.5390.8240.7030.552UNet++0.8360.7210.5700.8450.7320.586SegResNet0.7910.6650.4870.8080.6820.516Swin-UNETR0.8270.7200.5620.8430.7390.574UNETR0.8140.6930.5310.8390.7170.549CoT-UNet0.8300.7140.5680.8380.7250.576基于Transformer的3D U-Net0.8100.6720.5260.8370.6900.545ME-UNet0.8470.7280.5830.8620.7540.605表5. 计算效率和资源消耗的比较。方法每个周期的时间(秒)参数(M)GPU内存使用量(GB)3D U-Net29.098.2820.1UNet++18.0631.0618.4SegResNet26.28278.9518.3Swin-UNETR38.341477.4722.7UNETR13.61155.8411.8CoT-UNet4.8413.356.1基于Transformer的3D U-Net5.1911.885.1ME-UNet5.4517.914.7下载:下载高分辨率图像(124KB)下载:下载全尺寸图像(a)。AUC。下载:下载高分辨率图像(131KB)下载:下载全尺寸图像(b)。DSC。下载:下载高分辨率图像(123KB)下载:下载全尺寸图像(c)。IoU。图4. 不同模型的AUC、DSC和IoU随GPU内存使用的变化趋势。下载:下载高分辨率图像(129KB)下载:下载全尺寸图像(a)。AUC。下载:下载高分辨率图像(124KB)下载:下载全尺寸图像(b)。DSC。下载:下载高分辨率图像(125KB)下载:下载全尺寸图像(c)。IoU。图5. 所有模型在训练迭代过程中AUC、DSC和IoU值的比较。下载:下载高分辨率图像(128KB)下载:下载全尺寸图像图6. ImageCAS数据集上所有比较模型的ROC曲线。表6. ME-UNet不同架构变体的消融研究结果。这里,“Direct”表示在通道维度上直接将特征图与重建图像拼接。“Convolutional”表示用两个3×3×3卷积层替换USC,并进行实例归一化和ReLU激活。“USC”表示使用不确定性跳跃连接,“Channel”和“Spatial”分别表示仅应用通道级和空间注意力组件。组DirectConvolutionalUSCChannelSpatialAUCDSCIoUGPU内存使用量(GB)10.5320.1490.0093.12?0.7860.6770.5113.63?0.8160.6810.5164.04?0.8240.7090.5494.65??0.5810.1610.0153.16???0.8130.6870.5243.77???0.8280.7130.5524.18??0.8350.7180.5704.69??0.8330.7090.5594.7值得注意的是,可学习参数的数量较少并不一定意味着较低的计算成本或减少的GPU内存消耗。实际上,训练时间和内存使用主要受中间特征图的存储、数据移动以及涉及的操作类型的影响,而不仅仅是参数数量。特别是,具有复杂注意力机制或密集跳跃连接的模型尽管参数较少,但可能会产生较大的内存开销和更长的运行时间。相比之下,ME-UNet主要依赖于GPU高效的卷积操作和线性复杂度的注意力机制,即使在参数数量略多的情况下,也能实现更快的训练速度和更低的内存使用量。为了进一步证明所提出方法的内存效率,我们调整了模型的参数,以探索其在不同内存使用场景下的性能。图4显示了三个评估指标随内存使用量的变化曲线。通过调整参数,我们为每个模型获得了十个数据点。由于参数规模庞大且模型结构复杂,我们无法精确控制内存使用量,导致数据点分布不均。尽管如此,我们仍能观察到模型性能的变化趋势。随着GPU使用量的增加,性能指标的快速提升和早期平台期表明ME-UNet具有很高的效率。它能够在不需要大量计算资源的情况下达到接近最优的性能。这使得ME-UNet特别适合于GPU资源有限的现实世界医疗应用。此外,可以观察到性能提升迅速的模型(如ME-UNet、COT-UNet和基于Transformer的3D U-Net)在达到峰值后指标略有下降。这可能是由于过拟合造成的。
下载:下载高分辨率图像(2MB)
下载:下载全尺寸图像
图7. 不同模型在ImageCAS数据集的三个代表性样本上进行血管分割结果的视觉比较。
下载:下载高分辨率图像(1MB)
下载:下载全尺寸图像
图8. 两个代表性案例的消融研究结果的视觉比较。
为了进一步说明所提出方法的性能,图5显示了所有模型在测试数据上AUC、DSC和IoU随迭代次数的变化情况,每十次迭代显示一个点。对于AUC,可以观察到所有模型最初都有所提升,然后达到峰值,随后逐渐下降。CoT-UNet在40次迭代时达到所有模型中最高的AUC;然而,在后续迭代中下降,并最终低于ME-UNet、UNet++和基于Transformer的3D U-Net。ME-UNet在大多数迭代中保持了较高的AUC。关于DSC,ME-UNet在70次迭代时达到接近0.75的最高值。尽管在峰值后略有下降,但在后续迭代中仍保持所有模型中最高的DSC。对于IoU,ME-UNet也在70次迭代时达到约0.6的最高值。这些结果表明ME-UNet在不同迭代中始终保持较高的性能。此外,我们比较了不同模型的ROC曲线,如图6所示。显然,ME-UNet的ROC曲线优于其他模型,表明其分类能力更强。
4.4. 消融研究
为了验证ME-UNet中每个模块的有效性,我们在ImageCAS数据集上进行了消融研究,结果如表6所示。第1组的结果表明,在移除跳跃连接和注意力模块后,性能显著下降,同时显示出最低的GPU内存消耗为3.1 GB。在第5组中,移除了跳跃连接但保留了DFA模块,观察到比第1组略有改进,AUC、DSC和IoU分别增加了0.049、0.012和0.006,同时保持了相同的内存使用量。这表明在没有跳跃连接的情况下,仅使用注意力机制的效果有限。在第2、3和4组中引入跳跃连接后,性能得到了显著提升,但GPU内存使用量也增加了。在不同的跳跃连接设计中,USC(第4组)的分割精度优于直接跳跃(第2组)和卷积跳跃(第3组),AUC分别提高了0.038和0.008,DSC分别提高了0.032和0.028,IoU分别提高了0.038和0.033。尽管USC的内存成本为4.6 GB,比简单的跳跃连接更高,但它仍然比传统的基于注意力的跳跃设计更高效,显示出精度和内存消耗之间的良好平衡。将第2组(仅直接跳跃)与第6组(直接跳跃加上DFA)进行比较,添加DFA后,AUC、DSC和IoU分别提高了0.027、0.010和0.013,而GPU内存使用量仅略有增加。同样,比较第3组和第7组,加入DFA后性能持续提升,但内存使用量从4.0 GB增加到4.1 GB。这些结果表明DFA在最小化内存开销的同时增强了特征表示。此外,从完整的ME-UNet配置中移除空间注意力(第8组)或通道注意力(第9组)会导致性能明显下降,尽管GPU内存使用量相当。这表明通道注意力和空间注意力都为分割任务提供了互补的信息。总体而言,完整的ME-UNet实现了最佳的性能-效率平衡,提供了卓越的分割精度和适度的GPU内存消耗。
4.5. 案例研究
为了比较不同模型的性能,图7展示了ImageCAS数据集上三个样本的血管分割视觉结果。在案例1中,所有模型中的黑色虚线框内的血管都显得不连续。Swin-UNETR和UNETR显示了许多断裂的血管,而ME-UNet显示的血管断裂情况较少。在黄色虚线框中,ME-UNet成功识别出了血管段,而其他模型则出现了显著的血管段丢失。在绿色虚线框中,所有模型都错误地检测到了实际上不存在的血管,这可能是由于过拟合造成的。在案例2中,SegResNet、UNETR、CoT-UNet和基于Transformer的3D U-Net都错过了黑色虚线框内的血管段。在黄色虚线框中,只有ME-UNet正确识别出了更完整的血管段,而其他模型则出现了明显的血管丢失。在绿色虚线框中,所有模型都错过了血管段,其中SegResNet、Swin-UNETR和UNETR显示了血管不连续。在案例3中,所有模型在黑色虚线框内都显示出明显的血管不连续,尽管3D U-Net、UNet++和CoT-UNet的表现稍好一些。在黄色虚线框中,ME-UNet成功识别出了血管段,而其他模型则出现了不同程度的血管丢失。总体而言,ME-UNet在处理血管不连续性和最小化细节丢失方面优于其他模型。
表7. 训练过程中不同跳跃级别下体素级不确定性(σ)的统计特征。这里,Mean和Std表示σ的平均值和标准差,而P50、P90和P99表示不确定性分布的50%、90%和99%分位数。
时代 跳跃次数 平均值 标准差 P50 P90 P99
1 10.007 0.0006 0.0068 0.0081
2 0.0082 0.0017 0.0075 0.0100 0.0144
3 0.0083 0.0036 0.0072 0.0100 0.0266
4 0.0106 0.0058 0.0085 0.0171 0.0423
5 0.0097 0.0020 0.0102 0.0138
6 0.0387 0.0489 0.0186 0.0897 0.2576
7 0.0438 0.0449 0.0237 0.1137 0.1647
8 0.2023 0.1940 0.1787 0.5002 0.8118
9 0.0079 0.0034 0.0089 0.0110 0.0188
10 0.1233 0.1290 0.0628 0.3394 0.4504
11 0.8654 0.4961 0.8185 1.5709 1.8350
12 0.5518 0.5443 0.2970 1.4443 2.0299
另一方面,我们在图8中展示了所有消融组的视觉结果。在案例4中,第1组和第5组的性能较差,主要是由于移除了跳跃连接的负面影响。比较第2组和第6组在三个框中的表现,第6组通过识别更多的血管段而优于第2组。比较第3组和第7组以及第4组和ME-UNet时也观察到了类似的现象,表明DFA在捕捉血管结构方面是有效的。然而,在第8组和第9组中,黄色虚线框中的段丢失情况比第4组更严重。这表明空间注意力和通道注意力的组合在提高性能方面比单独使用更有效。此外,第3组和第7组在所有三个框中的性能都优于第2组和第6组,表明USC的表现优于卷积跳跃连接。在案例5中也可以观察到类似的趋势。这些结果突显了DFA和USC的有效性。
4.6. USC分析
为了进一步研究所提出的USC的行为,我们分析了训练过程中不同跳跃级别下估计的体素级不确定性的统计特征,如表7所示。不确定性值显示出明显的阶段性和层次性模式,而不是随机或均匀分布的。在训练的早期阶段,不确定性较低且在所有跳跃连接中相对均匀,表明模型尚未学会区分可靠和不可靠的特征区域。随着训练的进行,不确定性在跳跃级别之间变得越来越明显,深度跳跃连接显示出更高的方差和明显的长尾分布。在收敛时,浅层跳跃连接始终保持较低的不确定性,反映了低级空间特征的稳定性,而深层跳跃连接将更高的不确定性集中在一小部分体素中。这种依赖于层次和长尾的行为表明USC学习了与特征可靠性和语义模糊性相关的有意义的不确定性,为跳跃连接级别的不确定性感知调制提供了实证支持。
5. 结论
在本文中,我们提出了ME-UNet,这是一种用于CCTA中冠状动脉分割的内存高效3D U-Net架构。通过引入USC和DFA模块,ME-UNet在跳跃级融合过程中抑制了噪声特征的传播,同时以较低的计算开销捕获了全局结构信息。广泛的实验表明,ME-UNet在需要的内存远低于最先进模型(UNet++)的情况下,实现了卓越的分割精度,将内存使用量减少到了UNet++所需的大约25%。
准确的冠状动脉分割是下游临床任务(如冠状动脉中心线提取、斑块和狭窄评估以及介入计划)的关键前提。通过显著降低GPU内存需求同时保持高分割性能,ME-UNet促进了自动化冠状动脉分析流程在常规临床工作流程中的集成。从实际部署的角度来看,ME-UNet的低内存占用和竞争性的推理速度使其能够在常见的单GPU医院工作站或共享服务器上部署。此外,USC产生的不确定性估计提供了可解释的线索,突出了模糊的血管边界和小分支,支持了质量控制和临床医生与AI的合作。
未来的工作将致力于为USC和DFA建立更严格的理论基础。具体来说,我们旨在从特征可靠性和信息传播的角度分析不确定性引导的跳跃连接,并进一步研究DFA作为一种内存感知的全局上下文聚合机制,以提高内存高效3D分割模型的可解释性、泛化和临床适用性。
CRediT作者贡献声明
Quan Haodi:撰写——原始草稿、可视化、软件、方法论、概念化。
Liatsis Panos:撰写——审阅与编辑、撰写——原始草稿、监督、项目管理、资金获取、数据管理。
伦理声明
本研究遵循了机构和国家研究委员会的伦理标准以及1964年赫尔辛基宣言及其后续修订案。研究使用了从公开可用的ImageCAS数据集中获得的回顾性、匿名化的冠状动脉CT血管造影(CCTA)图像,这些图像是在原始数据提供者的适当伦理批准和知情同意下收集的。本研究没有获取涉及人类参与者或动物的新数据,所有分析都是在去标识化的数据集上进行的。因此,本研究不需要额外的伦理批准或参与者的知情同意。