NVIDIA 3090 实现 TransGUNet | 注意力机制下跨尺度图与熵驱动特征选择(EFS),提升图像分割性能 !

技术

picture.image

跳层连接工程主要用来解决编码器和解码器之间的语义差距,同时整合全局依赖性以理解医学图像分割中复杂解剖结构之间的关系。虽然已有模型提出了基于Transformer的方法来在跳层连接中引入全局依赖性,但它们往往在捕捉细节局部特征时面临高计算复杂度的限制。相比之下,图神经网络(GNNs)利用图结构有效地捕获局部和全局特征。

利用这些特性,作者引入了一种注意机制下的跨尺度图神经网络(ACS-GNN),通过将跨尺度特征图转换为图结构,并利用节点注意机制捕获复杂的解剖结构,从而增强跳层连接框架。

此外,作者观察到深度学习模型往往会产生无信息特征图,这会降低空间注意力图的质量。为此,作者结合了基于熵驱动的特征选择(EFS)与空间注意力机制,计算每个通道的熵分值并筛选掉高熵特征图。

作者的创新框架TransGUNet由ACS-GNN和基于EFS的空间注意力组成,通过利用GNNs以及可靠的空间注意力图,有效增强了跨模态领域的泛化能力,确保跳层连接中的更稳健特征。

通过全面的实验和分析,TransGUNet在六个已见数据集和八个未见数据集上的分割性能表现优异,相对于先前方法展现了显著更高的效率。

  1. Introduction

医学图像分割对于早期发现异常组织和制定治疗计划至关重要[9]。传统分割算法受到了医学专家的高度关注[23,32,47,57]。然而,由于医学图像中存在的严重噪声、非均匀强度分布以及各种临床环境,这些方法仍然缺乏普遍适用性[51]。因此,这一问题引发了对基于计算机的诊断程序可靠性的担忧[18]。

最近,卷积神经网络(CNNs)因其在捕捉局部和空间层次特征方面的稳健性而在医学图像分割中得到了广泛的应用[44, 52, 77]。然而,基于CNN的模型在捕捉理解医学图像中复杂解剖结构所需的全局依赖方面存在困难[24]。这一局限性促使研究行人开始使用Transformer提取全局依赖以进行医学图像分割[5, 7]。尽管Transformer模型具有许多优势,但在编码器和解码器之间仍难以弥合语义差距,这限制了它们充分利用全局依赖性的能力,并导致分割性能欠佳[64]。

为了减少医学图像分割中的语义差距,已经积极采用了多种模型来改进 Short-Cut 结构。其中最具有代表性的尝试之一是

[79],它通过密集连接在 Short-Cut 中实现了跨尺度特征融合。类似地,UCTransNet [64] 和 CFATransUNet [63] 则采用基于 Transformer 的方法,从通道视角捕捉特征图的局部跨通道交互。然而,这些模型由于其密集连接、大量使用 Transformer 块以及医学图像中存在的严重噪声而导致计算复杂度增加和注意力不明确。因此,在不引入模糊性的情况下高效地利用全局依赖性并减小编码器和解码器之间的语义差距,对于克服这些挑战和提高医学图像分割性能至关重要。

为回答这个问题,作者关注了图神经网络(GNNs),这种网络特别适合灵活而有效地捕捉局部和全局依赖关系,使其成为复杂视觉感知任务的理想选择[21]。通过利用这一能力,作者提出了一种注意力驱动的跨尺度GNN(ACS-GNN),它能够高效地缩小编码器和解码器之间的语义差距。它将跨尺度特征图转换为图,并对每个节点应用注意力机制,以促进稳健的特征整合。此外,作者观察到深度学习模型经常产生无信息性的特征图,这会降低空间注意力图的质量[8, 53]。为了解决这一问题,作者引入了一种基于熵的特征选择方法(EFS),该方法计算每个通道的熵并滤除高熵通道。通过结合ACS-GNN和EFS驱动的空间注意力机制,作者设计了一种新的医学图像分割模型——TransGUNet,它可以有效地捕捉 Patch 之间的关系,而不受病灶大小和 Patch 之间距离的影响(图1)。广泛的实验结果表明,作者的基于图的方法始终优于基于 Transformer 和卷积的方法。因此,TransGUNet代表了医学图像分割领域 Shortcut 框架的一个重要进展,并提供了一种稳健且高效的解决方案来应对现有挑战。本研究的主要贡献可以总结如下:

picture.image

作者提出了一种名为TransGUNet的新颖医疗图像分割模型,该模型利用了基于跨尺度图神经网络(GNN)的 Shortcut 框架,且不含有歧义的空间注意力机制,并适用于各种模态和临床场景。据作者所知,作者的新型 Shortcut 框架是首次成功且有效地利用具有明确空间注意力的注意力跨尺度GNN进行医疗图像分割的研究。

所提出的注意力驱动的跨尺度图神经网络(ACS-GNN)使模型能够理解医疗图像中的复杂解剖结构。此外,作者还结合了基于熵的特征选择(EFS)与空间注意力,以生成更可靠的空間注意力图。

实验结果表明,与用于各种临床场景的基于 Transformer 和卷积的方法相比,TransGUNet显著优于现有方法在医疗图像分割上的表现。

  1. Related Works

医学图像分割中的 Shortcut 工程。UNet引入 Shortcut 标志着一个重要的里程碑,并使其成为医学图像分割中应用最广泛的 Baseline 模型。然而,编码器和解码器之间仍然存在语义差距,导致性能不佳 [38]。这一问题驱动了近期对 Shortcut 进行改进的努力,以尽量减小这种语义差距。

是一种代表性模型,它结合了密集连接和邻居尺度特征的 Shortcut 。此外,MSNet [77] 和 M2SNet [78] 通过使用减法模块减少冗余特征来设计更高效的模型。

最近,transformer 被用作 Shortcut 模块,以捕捉医学图像中的全局依赖性 [17, 26]。值得注意的是,UCTransNet [64]、FCT [58] 和 CFATransUNet [63] 通过利用基于transformer的 Shortcut 框架保持了全局依赖性。然而,这些模型具有复杂的结构,参数量超过60M,计算成本高且效率低。作者的创新性TransGUNet通过利用跨尺度图神经网络及节点注意力机制并采用更为高效的小于25M参数量的架构解决了这些问题。作者在表1和附录(图8)中比较了各种 Shortcut 框架的方案和特性。

picture.image

计算机视觉中的GNN应用。传统上,GNN被用于自然语言处理[33]和推荐系统[73],这是因为它们能够理解数据集内部错综复杂的关系。近年来,在计算机视觉领域,GNNs [21, 22] 被积极研究,以灵活地基于图提取全局依赖和局部特征,而图是一种泛化的数据结构,涵盖了网格(CNN)和序列(Transformer)。

例如,SFDGNet [68] 使用GNN提取内容特异性的人工频率特征,并理解复杂的空间和频率关系。另外,GazeGNN [62] 将原始的眼球追踪数据与图像集成到统一表示的图中,用于实时疾病分类。特别是在医学图像分割方面,ViGUNet [29] 和 PVTGCASCADE [50] 利用GNN处理复杂的解剖结构。此外,GTBA-Net [72]、TSGCNet [12]、MSAGAANet [69]、TGNet [75] 和 GSENet [36] 尝试将transformer与GNN结合应用于医学图像分割。然而,这些模型并未考虑跨尺度特征,缺乏可靠的空域注意力图。鉴于这些局限性,作者精心设计了TransGUNet,该模型通过ACS-GNN结合EFS实现跨尺度特征融合,并充分利用全局依赖关系。

  1. Method

3.1. Encoder and Decoder in TransGUNet

作者使用了一种包含多个基于池化结构的多头自注意力机制的金字塔池化Transformer(P2T)[70]。P2T相比视觉Transformer(ViT)和金字塔视觉Transformer(PVT),在计算成本显著降低的同时,具有更高的表示能力,并被广泛应用于各种医学影像分割模型中[10, 37, 76]。受先前研究的启发,作者将解码器的编码架构与解码器相同,以充分利用全局依赖关系。尽管作者在实验中主要使用P2T进行展示,作者还在附录(表9)中提供了多种CNN和Transformer Backbone 网,以展示所提出方法在不同 Backbone 网架构上的灵活性和稳健性。

picture.image

3.2. ACS-GNN with EFS-based spatial attention for Skip Connection

动机:人的视觉系统(HVS)通过将物体划分成大部件并基于每个部件的连接强度来识别物体 [39]。这一过程有助于通过识别图像不同部分之间的关系来解释复杂的场景,从而获得对物体及其相互作用的整体理解 [41, 48]。受到这些原则的启发,作者的方法采用类似HVS的策略,即将跨尺度特征图转换为图形,以在高维特征空间中理解和分析复杂的解剖结构。

然而,显著的噪声和复杂的背景会产生高度模糊的视觉信号,干扰神经系统的处理。HVS通过信号过滤和注意力处理来缓解这一问题 [49, 59]。因此,作者提出了一种基于熵的特征选择策略,该策略模拟了这些特征过滤和注意力处理过程,称为基于EFS的空间注意力。这些组件的集成增强了注意力机制在保留全局依赖性和局部细节方面的表现,而不会造成模糊。

TransGUNet的整体架构如图2所示。ACS-GNN与基于EFS的空间注意力可以分为四个步骤:1)特征预处理,2)ACSGNN(图2(c)),3)基于EFS的空间注意力(图3),4)特征后处理。

picture.image

picture.image

特征预处理。设

为第

层编码阶段的特征图,其中

是输入图像的分辨率。由于每个阶段的通道数主要影响解码器的复杂性,作者使用了大小为

的 2D 卷积来将通道数减少到

。为了获得跨尺度特征图

,作者将这些特征图在

时调整到相同的分辨率,具体操作如下:

其中,

和 Resize

分别表示具有核大小

的 2D 卷积操作和一种将特征图调整为目标空间分辨率

的操作。如果输入特征图的目标分辨率与原分辨率不同,则使用双线性插值将其上采样或下采样以匹配目标分辨率。如果两者有相同的分辨率,则不需要进行重新调整尺寸,如图2(b) 所示的 “分辨率固定” 操作。随后,作者将每个重新调整尺寸后的特征图进行拼接,

,其中

表示沿通道维度对特征图进行拼接。为了方便起见,作者假设

。然后,作者应用注意力跨尺度图神经网络。在获得跨尺度特征

后,作者将进一步进行处理。

使用核大小为

的二维卷积和批量规范化(BatchNorm)对特征图进行处理。随后,该特征图被展平为向量

,其中

。向量

中的每个像素点充当图中的节点。此外,在每个展平的向量元素中添加了一个相对位置向量以保留位置信息。接着,作者使用膨胀

最近邻(KNN)算法构建了特征图。为了实现节点间的信息交换,作者采用了Max-Relative图卷积(MRConv)[34],因为它简单且高效,无需学习节点聚合的可学习参数。

图卷积的过程通过以下方式实现:

。为了通过自适应加权改进特征聚合,作者应用了节点注意力机制,以便在学习关键特征的相关性的同时优先考虑这些特征。基于ECANet[65],作者采用单一的一维卷积操作以及核大小为

的Sigmoid函数设计了节点注意力机制。首先,使用全局平均池化(Global Average Pooling)和全局最大池化(Global Max Pooling)分别将

压缩为

,然后将每个统计量聚合生成一个节点注意力图。最后,这种注意力机制可以方便地实现如下:

其中,

分别表示核大小为

的 1D 卷积和 Sigmoid 函数。作者将展平向量

重新reshape 成原始特征图形状,并应用核大小为

的 2D 卷积和BN以增加非线性性,从而获得精炼特征图

。为了应对GNN中的过度平滑问题[35, 46],作者还利用了具有两个连续卷积层和残差连接的 FFN ,如下所示:

其中,

分别表示用于非线性的ReLU激活函数和批量规范化。具有熵驱动特征选择的空间注意力机制。尽管GNN可以维护全局依赖性,但在医学图像中存在的各种噪声和复杂背景仍然会导致包含高熵分数的无信息特征图,从而产生效果较差的空间注意力图[8, 53]。为了解决这一问题,作者提出了一种基于熵驱动特征选择(EFS)的空间注意力机制,如图3所示,该机制能够筛选出无信息特征图并生成一个更为可靠的空间注意力图。作者观察到,在高熵通道中,特征在像素 Level 上的激活通常是均匀或嘈杂的,如图4所示。因此,借鉴先前的研究[61, 66],作者使用香农熵计算输入特征图每个通道的像素 Level 的熵得分

,如下所示:

picture.image

随后,作者仅保留了具有最低

个熵得分的特征图通道,并使用这些通道进行空间注意力处理:

其中,排序操作

表示将元素按照升序排列。因此,

表示特征图

中低熵分数的Bottom-

索引。作者使用了 Introselect 算法,这是 PyTorch 默认的排序算法。ACS-GNN 和 EFS 基础的空间注意力机制的双重过程使作者的模型能够理解医学图像中的复杂解剖结构并生成更可靠的空域注意力图。

特征后处理。作者首先沿着通道维度将

分割为等数量的通道

。随后,每个特征图被重新调整大小,并与原始特征图之间应用残差连接,以增强训练稳定性,其中

的操作如下:

3.3. Training Procedure

作者采用带有深层监督的多任务学习来增强模型的表示能力并缓解梯度消失问题。为了实现这一点,作者在每个阶段应用核大小为

的二维卷积和 sigmoi d 函数,并结合上采样,从

中获得了四个预测

,其中

。作者将

表示为分别代表区域

和边界

的多个输出。总损失函数定义为:

其中,

分别表示区域和边界的 Ground Truth。作者通过应用各向异性 Sobel 边缘检测滤波器 [31] 获得了

。区域预测的损失函数定义为

,其中

分别是加权 IoU 和二元交叉熵(BCE)损失函数。该损失函数在之前的研究所中已被一致定义 [15, 45, 78]。此外,作者定义了边界损失函数

为 BCE 损失函数。

  1. Experiment Results

4.1. Experiment Settings

每种模型都在五个医学分割任务上进行了训练和评估,包括多器官、皮肤癌、COVID-19感染、乳腺肿瘤和息肉。为了方便起见,作者将已见临床设置(表2)作为测试数据集,其分布与训练数据集相同。此外,作者还在八个外部分割数据集上评估了每个模型在不同分布条件下的领域泛化能力,这些不同的分布被称作未见临床设置(表3)。

由于篇幅限制,详细的数据库描述和拆分信息详见附录(表6和表7)。为了评估每种模型的性能,作者选择了两个指标:Dice分数系数(DSC)和平均交并比(mIoU),这两个指标在医学图像分割中广泛使用。此外,附录中还提供了更多不同指标的定量结果。

picture.image

picture.image

picture.image

picture.image

作者比较了所提出的TransGUNet(作者的方法)与十二个代表性的医学图像分割模型,包括UNet [52]、

[79]、CENet [19]、TransUNet [7]、MSRFNet [55]、DCSAUNet [71]、M2SNet [78]、ViGUNet [29]、PVT-GCAS [50]、CFATUNet [63]、MADGNet [45]以及GSENet [36]。

在所有结果中,作者报告了三次试验的平均性能以确保可靠性。在所有表中,粗体和斜体分别表示第一和第二的最佳性能结果。并且,在表2和表3的最后一行中,指出了TransGUNet与其他第二最佳方法之间的性能差距。

4.2. Implementation Details

作者使用Pytorch 1.8在单块NVIDIA RTX 3090上实现了TransGUNet。

多器官分割。作者借鉴了之前的研究[63],使用Adam优化器和学习率0.001进行多器官分割。作者在批量大小为24的情况下优化每个模型,并训练150个周期。在训练过程中,以50%的概率进行了翻转操作,并对图像进行了-20°到20°之间的旋转处理。由于在已知和未知临床环境中使用的体积具有不同的分辨率,所有图像均被调整为224×224的尺寸。值得注意的是,在CFATUNet中,这是最近的一种多器官分割模型,作者也采用了相同的设置。作者希望通过这些一致的设置来进行公平的比较,从而训练所有模型。

二进制分割。作者使用Adam优化器初始学习率为

,并通过余弦退火学习率调度器将每个模型的参数减小到

。作者采用批量大小为16对每个模型进行优化,并在结肠息肉、皮肤癌、乳腺肿瘤和COVID-19感染分割任务中分别训练50、100、100和200个周期。在训练过程中,作者以50%的概率应用水平和垂直翻转,并且包括从

的角度旋转,这是多尺度训练策略的一部分。这种方法常用于医学图像分割模型[15, 45, 77, 78]。由于每个数据集中的图像分辨率不同,所有图像都被调整为

。此外,作者还需要澄清,在M2SNet和MADGNet中,作者使用了相同的设置,这两种方法是最具代表性的医学图像分割方法。

4.3. Comparison with State-of-the-art models

作者使用了表2中所示的相同模型来评估其在未见过的临床设置中的领域泛化能力。

为了方便起见,作者用

表示TransGUNet与其他模型在已知和未知临床场景下性能提升的差距。根据表2和表3所示,TransGUNet在平均各数据集上的分割性能最高。与专注于增强 Shortcut 框架的UNet++、M2SNet和CFATransUNet相比,TransGUNet分别展示了DSC改进百分比为

。此外,与使用单尺度GNN并带有空间注意力机制的PVT-GCASCADE相比,TransGUNet展示了DSC改进百分比为

。尽管MADGNet在已知临床场景中达到了最先进的性能,但TransGUNet在未知临床场景中平均展示了高达

的显著DSC提升。令人惊讶的是,只有当TransGUNet在Synapse(CT模态)上进行训练时,在AMOS-MRI(MRI模态)上才能实现超过

的DSC。

这些结果表明,采用具有明确空间注意力机制的注意力跨尺度GNN为基础的 Shortcut 对于理解医学图像中的复杂解剖结构至关重要。图5表明,TransGUNet包含约25.0M个参数,浮点运算次数(FLOPs)为10.0G,这在计算效率方面具有明显优势。作者将在附录(表8)中详细列出每个模型的参数数量、FLOPs以及推理时间(ms)。

picture.image

picture.image

图6展示了各种模型的定性结果。UNet、CENet、MSRFNet和DCSAUNet未包含全局依赖关系或减少编码器和解码器之间的语义差距,因此产生的预测结果嘈杂且不可靠。尽管UNet+ 和M2SNet通过嵌套卷积和减法单元分别减少了语义差距,但它们在包含复杂息肉结构的结肠镜检查图像或含有严重噪声的超声图像中的预测仍然不可靠。

虽然TransUNet利用了全局依赖关系,但它并未在解码器中使用,导致预测不准确。CFATransUNet通过减少语义差距并在解码阶段引入Transformer块改进了先前的方法,但忽略了关键的空间关系,未能捕捉到有效解释医学影像所需的精细和局部细节。

尽管PVT-GCASCADE在一定程度上缓解了这些缺陷,但它未能考虑跨尺度交互、可靠的空域注意力以及图中每个节点的重要性。此外,如图1所示,TransGUNet可以理解语义相似 Patch 之间的关系,不论它们之间的距离如何。

因此,尽管存在严重噪声、不同大小的病灶和多种模态下的复杂解剖结构,TransGUNet仍能产生可靠的预测结果,这得益于ACS-GNN和EFS基空域注意力的双重利用。

picture.image

picture.image

4.4. Ablation Study on TransGUNet

作者在结肠息肉分割任务上进行了消融研究,以展示ACS-GNN和EFS基空间注意力的有效性。作者希望澄清,在消融研究中使用的实验设置与主要实验相同,以便进行公平比较。

ACS-GNN消融研究。如表4所示,作者的方法(S4)在所有设置中表现出最佳性能。SO表示不包含ACS-GNN的TransGUNet。最显著的结果是,在GNN中应用 Shortcut ,无论是单尺度(S1)还是跨尺度(S3),都能改善已见过数据集上的性能。然而,在未见过的数据集上,单尺度GNN的性能下降了0.4%,而跨尺度GNN的性能则提高了0.7%。这些实验结果解释了为什么TransGUNet能够优于PVT-GCASCADE。

此外,NA改善了单尺度(S2)和跨尺度(S4)GNN的性能,由于采用了类似于ECA的注意力机制,参数量和FLOPs几乎没有增加。因此,采用ACS-GNN的TransGUNet在多种模态和临床设置下表现明显更好。

picture.image

基于EFS的空间注意力模块的消融研究。如实现细节部分所述,由于作者设定了(C_{r}=64), Short-Cut 中的总通道数等于(4C_{r}=256)。因此,当(4C_{r}=M),即选定的通道数时,所有特征图均用于生成空间注意力图。

图7中,((\cdot,\cdot))表示在息肉分割的已见和未见临床设置上的DsC。图7显示,(M=256)的结果比使用EFS为基础的空间注意力模块的几种情况((M={8,16,32,64,128}))都要差。这些结果表明,作者的特征选择方法能够提高空间注意力的质量,这是此前尚未解决的问题。如果作者选择更少的通道数((M={8,16,32})),空间注意力图的不确定性会增加,因为缺乏信息。

picture.image

在ACS-GNN中关于目标分辨率的消融研究。在本节中,作者进行了一项消融研究,以比较不同目标分辨率

下的性能和效率。每个分辨率对应Transformer编码器提取的特征图

,其中

。作者将

作为目标分辨率

。如表5所示,实验结果显示,分辨率

取得了最佳性能。尽管分辨率

也表现出色,但由于输入特征图分辨率提高导致计算量增加以及GNN固有的特性,其效率显著下降。这表明虽然更高的分辨率能够捕获更多的细节信息,但同时也需要付出更多的计算资源代价。

picture.image

  1. Conclusion

在本研究中,作者提出了TransGUNet,这是一种创新的医学图像分割模型,将ACSGNN和基于EFS的空间注意力相结合。

该提出的模型通过减少编码器和解码器之间的语义差距、保留关键信息而不引入模糊的空间注意力,并利用全局依赖性,有效地解决了现有模型的局限性。

通过在六个已见数据集和八个未见数据集上的广泛实验,TransGUNet展现了优于最新模型的优越性能和更高的效率。消融研究表明,将GNNs纳入 Short-Cut 工程显著增强了模型捕获和利用复杂解剖特征的能力。

此外,EFS确保只考虑最有信息量的特征,从而提高了空间注意力图的质量。因此,TransGUNet代表了医学图像分割的一个重要进步,提供了一个稳健、高效且准确的模型,可以在各种医学应用中使用。

在未来的工作中,作者将专注于优化内存效率,并探索其在实际医疗保健环境中的部署。

参考

[0]. TransGUNet: Transformer Meets Graph-based Skip Connection for Medical Image Segmentation .

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
基于火山引擎 EMR 构建企业级数据湖仓
火山引擎 EMR 是一款云原生开源大数据平台,提供主流的开源大数据引擎,加持了字节跳动内部的优化、海量数据处理的最佳实践。本次演讲将为大家介绍火山引擎 EMR 的架构及核心特性,如何基于开源架构构建企业级数据湖仓,同时向大家介绍火山 EMR 产品的未来规划。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论