超越传统 UNet ，GCtx-UNet 结合全局与局部特征，实现高效图像分割 ！

医学图像分割对于疾病诊断和监测至关重要。尽管当前的分割网络如UNet在效果上有效，但它们在捕捉远程特征方面存在困难。更准确的模型如TransUNet、Swin-UNet和CS-UNet具有更高的计算复杂性。 为了解决这个问题，作者提出了GCtx-UNet，一个轻量级的分割架构，它可以捕获全局和局部图像特征，其准确度优于或与现有技术水平相当。GCtx-UNet使用视觉 Transformer ，该 Transformer 利用全局上下文自注意力模块结合局部自注意力来建模长距离和短距离的空间依赖关系。 GCtx-UNet在Synapse多器官腹部CT数据集、ACDC心脏MRI数据集以及几个息肉分割数据集上进行了评估。 在Dice相似系数（DSC）和Hausdorff距离（HD）指标方面，GCtx-UNet优于基于CNN和基于Transformer的方法，尤其是在分割复杂和小型解剖结构方面取得了显著增益。此外，GCtx-UNet比现有技术水平的方法在模型大小、计算工作负载和训练推理速度上都要高效得多，使其成为临床应用的实际选择。

1 Introduction

自动化医学图像分割在提供有价值信息以预防、诊断、进展监测和预测各种疾病以及定量病理学评估中至关重要。目前，包括编码器、解码器和跳跃连接在内的U形深度神经网络在医学图像分割中应用最为广泛。尽管U形网络在许多医学图像分割任务中取得了最先进的表现，但仍然存在局限性。一个主要的局限性是编码器在有效提取和整合长距离和局部特征方面的能力。

基于卷积神经网络（CNN）的方法，如UNet [26]和UNet++ [35]，在捕捉局部特征方面表现出色，但它们在建模数据中的长距离依赖关系方面存在困难。而基于Transformer的方法，如Swin-UNet [6]，虽然能够建模长距离像素关系，但在建模局部信息时缺乏空间感应偏差，导致结果不理想。

以往的研究探索了如TransUnet [8]之类的CNN-Transformer混合架构，以捕捉全局和局部信息，但这些模型通常显著增加了参数数量。这进而导致计算复杂度增加，可能限制其实际应用。最近，Hatamizadeh等人[13]提出了一种全局上下文视觉Transformer（GC-ViT），它利用全局上下文自注意力模块，并与局部自注意力相结合，有效地和高效地建模长距离和短距离的空间交互。GC-ViT在图像分类、目标检测和语义分割任务上取得了最先进的结果。

在本文中，作者介绍了GCtx-UNet，这是一种为医学图像分割设计的类似UNet的分割网络。GCtx-UNet使用带有跳跃连接的GC-ViT [13]编码器和解码器有效地捕捉长距离和短距离的语义特征。

这种架构在提高性能的同时，需要的模型参数更少，具有更高的推理速度和更低的计算复杂度。

作者在几个医学图像数据集上评估了GCtx-UNet的分割和运行时性能，包括Synapse、ACDC和几个息肉图像数据集。作者的实验结果显示，GCtx-UNet的性能优于或可媲美包括基于CNN、基于Transformer和混合分割网络在内的最先进的分割算法。同时，GCtx-UNet具有最小的模型尺寸，使用的训练时间和推理时间最少。

此外，作者在ImageNet和MedNet上对GCtx-UNet进行了预训练，MedNet是从公开来源收集的20万张医学图像的一套数据。与在自然图像（即ImageNet）上进行预训练的GCtx-UNet相比，在领域内图像（即MedNet）上进行预训练的GCtx-UNet获得了更高的准确度。

2 Related work

基于卷积神经网络（CNN）的方法被广泛使用，并被认为是最突出的医学图像分割方法之一。特别是像UNet [26]及其衍生架构这种基于编码器-解码器的方法，在医学图像分割中表现出卓越的有效性。例如，Att UNet [25]通过注意力门增强了分割效果，而UNet++ [35]引入了一种替代的跳跃连接机制，即嵌套和密集型，在一定程度上减轻了UNet各层次之间的语义差距。这种改进相比于UNet带来了显著性能提升。然而，UNet++无法完全捕捉到全尺度的语义特征。黄等[15]提出了UNet3+，通过结合来自不同尺度的低级细节和高级语义，最大化利用全尺度特征图。基于CNN的方法已应用于各种医学图像分割任务，如在视网膜图像分割[11]和皮肤分割[33]中，展示了在实施和训练中的性能前景和实践性。基于ResNet架构的分割算法在医学图像分割[19]中确立了其地位。例如，Res-UNet [32]通过加权注意力机制增强了视网膜血管分割。

Transformer中固有的自注意力机制（MSA）使其能够执行全局相关性建模，有效处理长距离依赖。利用这种能力，Transformer在自然语言处理和计算机视觉任务中取得了显著进展，因其卓越的全局建模能力。一些开创性研究已经引入了基于Transformer的架构用于医学图像分割。曹等[6]提出了Swin-UNet，将Swin Transformer [21]整合到U形分割网络中用于多器官分割。阿扎德等[2]提出了TransDeepLab，用于皮肤病变分割，通过不同的窗口策略增强了DeepLab。此外，黄等[16]引入了MISSFormer，利用不同尺度上的全局信息进行心脏分割，而阿扎德等[3]提出了TransCeption，改进了 Patch 合并模块以在单一阶段捕捉多尺度表示。在编码器端结合卷积操作和Transformer，Transclaw UNet [7]能够进行详细分割和长距离关系学习。UNETR [12]采用序列到序列预测用于3D医学图像分割。这些发展凸显了基于Transformer的方法在医学图像分割中的变革性影响，为更广泛采用和深度学习进步指明了道路。

Swin Transformer [21]引入了局部窗口自注意力以减少计算成本，使其随图像大小线性增长，使用移位窗口注意力捕捉跨窗口信息，并通过分层架构利用多分辨率信息。然而，移位窗口注意力由于覆盖面积小，难以捕捉长距离信息，并且像ViT [9]那样缺乏归纳偏置。GC-ViT [13]是一个与Swin Transformer类似的分层架构，但使用全局窗口注意力而不是移位窗口注意力，有效捕捉长距离信息。GC-ViT还使用卷积层进行下采样，为网络提供像局部偏置和跨通道交互这样的理想属性，这些在ViT和Swin Transformer中是缺失的。GC-ViT有4个阶段，每个阶段都由局部和全局多头自注意力（MSA）层的交替块组成。如图1所示，在每个阶段，通过使用包含不同图像区域的全局上下文信息的全新融合倒残差块计算全局 Query 标记。而局部自注意力模块负责建模短距离信息，全局 Query 标记在所有全局自注意力模块之间共享，与局部键和值表示进行交互。

3 Pre-training on MedNet dataset

大多数基于CNN和Transformer的分割模型都是在如ImageNet这样的自然图像上进行预训练的。然而，这对于医学图像分割来说并不是最优选择，因为自然图像和医学图像模态之间存在语义差距[27, 1]。在这项工作中，作者预训练了GC-ViT[13]模型，具体是GCVit xxTiny，在一个名为_MedNet_的大型医学图像数据集上，该数据集包含超过20万张从几个公共数据集[30]和Kaggle[22, 23, 24]收集的医学图像。

MedNet包含不同类型的显微图像，如X射线、计算机断层扫描（CT）、光学相干断层扫描（OCT）和MRI。MedNet中的图像被分为65类。类似于Stuckner等人[29]和Alrfou等人[1]的方法，MedNet数据集被划分为训练集和验证集，每个类别在验证集中有100张图像，从而形成96.75%/3.25%的训练/验证分割。每个类别使用100张图像进行验证足以获得可靠的准确度指标，并在训练过程中防止过拟合。尽管验证集是平衡的，但训练集显示了一些类别不平衡。有些类别，每个类别包含的总图像不足0.12%。三个类别包含6.2%的图像。大多数类别拥有超过2000张图像，占训练集的1%到2%。MedNet包括来自X射线、CT、OCT和MRI等不同模态的图像，涵盖了广泛的医学疾病，如肾癌、宫颈癌、阿尔茨海默病、新冠肺炎、肺炎、肺结核、猴痘、乳腺癌和疟疾。

作者使用AdamW优化器[18]对GC-ViT xxTiny进行了100个周期的训练和测试，初始学习率为0.0001，权重衰减为0.05，并采用余弦衰减调度器。训练数据使用了albumentations库进行增强，包括随机改变对比度和亮度、垂直和水平翻转、光度失真和添加噪声。

训练过程一直持续到验证分数没有改善为止，采用早停准则，耐心度为10个周期。性能评估使用top-1和top-5准确度指标。Top-1准确度衡量的是正确标签为顶级预测的测试样本的百分比，而top-5准确度衡量的是正确标签出现在前五个预测中的测试样本的百分比。GC-ViT xx-Tiny模型的top-1准确度为82.3%，top-5准确度为98.2%。

4 Gctx-UNet Architecture

Gctx-UNet的核心组件是GC-ViT块，其局部和全局注意力机制如图1所示。如图2所示，每个GC-ViT块包括局部和全局的多头自注意力（MSA）、多层感知机（MLP）、全局 Token 生成器（GTG）和下采样层。GTG组件为计算添加全局上下文。局部MSA只能 Query 局部窗口内的块，而全局MSA可以在窗口内操作的同时 Query 不同的图像区域。在每一阶段，全局 Query 组件都是预计算的。该块还在下采样层中引入了基于CNN的模块，以包括归纳偏置，这是在ViT和Swin Transformer中缺失的对图像有用的特征。

Gctx-UNet是一个基于GC-ViT的U形编码器-解码器架构，具有用于长距离和短距离语义特征学习的跳跃连接。如图3所示，Gctx-UNet由编码器、瓶颈、解码器和跳跃连接组成。

1. 编码器和解码器都使用GC-ViT [13]来建模长距离和短距离的空间交互，无需计算注意力 Mask 或移动局部窗口等昂贵操作。
2. 在每个阶段，GC-ViT编码器和解码器由交替的局部和全局自注意力模块组成，以提取空间特征。两者都在像Swin Transformer那样的局部窗口中操作。
3. 跳跃连接将来自GC-ViT编码器的特征图与相应解码器阶段的特征图连接起来。瓶颈被用于获取深层特征表示，保持此组件中的特征维度和分辨率不变。
4. 编码器部分各阶段之间的下采样器和解码器部分各阶段之间的上采样器提供了理想的性质，如归纳偏置和建模通道间依赖关系。

在编码器中，初始图像被划分为四个块，作为四个阶段GC-ViT模块的输入。编码过程后，图像尺寸减小到(H/32) (W/32)。在解码器中，使用上采样操作将图像尺寸增加2倍，并将通道数减少2倍。通过跳跃连接将编码器每个阶段的特征与其在解码器中对应的阶段连接起来。解码器通过上采样来完成其任务。

Encoder

编码器采用分层GC-ViT方法来获取不同分辨率的特征表示。这是通过在4个阶段中减少空间维度同时将嵌入维度增加一倍来实现的。最初，输入图像 通过patchify层进行处理。这个层包括一个步长为2的 卷积操作以及填充，以生成重叠的图像块。随后，这些图像块通过另一个 卷积层投射到维度为 的嵌入空间中。在GC-ViT Backbone 网络的每个阶段之后，通过下采样层减少空间分辨率同时增加通道数。这个下采样操作有助于在不同分辨率下提取分层特征。

Downsampler

在降采样器中，作者采用了融合的MBConv模块来生成分层表示，通过向网络中注入归纳偏置，并建模通道间的相关性。然后，使用核大小为3、步长为2的卷积层来将空间特征分辨率降低2倍，同时将通道数翻倍。如图4所示，融合MBConv包括DW-Conv、GELU、SE和Conv。融合MBConv操作可以通过以下方程定义：

其中，DW-Conv指的是深度卷积，SE指的是挤压和激励块，GELU表示高斯误差线性单元函数。### 瓶颈

与Swin-UNet [6]类似，两个GC-ViT块用于构建瓶颈。瓶颈结构性地设计以促进深度特征表示的学习。在这个结构中，特征维度和分辨率保持不变。

Decoder

对称解码器与编码器相对应，采用GC-ViT Transformer块构建。解码器模仿了编码器的设计，将 Patch 块替换为非 Patch 块，嵌入层替换为去嵌入层，下采样块替换为上采样块。解码器的上采样块替代了编码器的下采样块。上采样块将相邻维度的特征图重构为更高分辨率的特征图，并将特征维度减半。这个上采样块有效地增加了空间分辨率，同时精炼和规范了特征表示，使其适用于在分割模型中解码和重建更高分辨率的特征。跳跃连接将编码器的特征与从上采样中恢复的深层特征融合在一起，因此减轻了由下采样产生空间数据损失。

5 实验评估

作者将GCtx-UNet的性能与包括基于CNN、基于Transformer和混合算法在内的最先进算法进行了比较。

表1总结了作者实验中使用的方法列表。作者评估了在MedNet和ImageNet上预训练的GCtx-UNet在三种不同类型的医疗图像数据集上的性能，包括Synapse多器官分割数据集（Synapse）、自动化心脏诊断挑战赛（ACDC)[5]，以及包括CVC-ClinicDB [4]、Kvasir-SEG [17]、CVC-300、ColonDB [31]和ETIS-LariDB [28]在内的息肉数据集。

Synapse: Synapse多器官分割数据集（Synapse）包括30个患者案例，共3779个轴向腹部临床CT图像，其中18个案例用于训练，12个案例用于测试。数据集包含8个腹部器官（主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏和胃）。每个CT体积包括个像素的图像切片， Voxel 空间分辨率为([)。ACDC: 自动心脏诊断挑战数据集（ACDC）[5]汇编了来自MICCAI 2017数据集的各种患者的MRI扫描结果。ACDC数据集包含100个心脏MRI扫描，每个扫描包含三个器官：右心室（RV）、心肌（Myo）和左心室（LV）。遵循TransUNet [8]，作者将数据集分为70个训练案例，10个验证案例和20个测试案例。Polyp: 作者使用了5个早期结直肠癌诊断图像的息肉数据集。CVC-ClinicDB [4]和Kvasir-SEG [17]数据集用于二值分割。CVC-ClinicDB数据集包含来自MICCAI 2015的612个带有标记息肉的RGB结肠镜图像，像素分辨率为。Kvasir-SEG数据集包含1000个息肉图像，像素分辨率从332487到19201072不等，以及相应的 GT 情况。遵循PraNet [10]中的设置，作者从CVC-ClinicDB数据集中使用了900张图像和Kvasir数据集中的548张图像进行训练。CVC-ClinicDB剩余的64张图像和Kvasir的100张图像用作测试集。为了评估泛化性能，作者在三个未见过的数据集上测试了模型：CVC-300、CVC-ColonDB和ETIS-LariDB。

Implementation

GCtx-UNet是使用PyTorch库实现的，并在配备12GB内存的Nvidia GeForce GTX TITAN X上进行训练。输入图像大小被缩小到224×224像素。作者的模型以批量大小24，学习率0.0001，使用带有动量0.9和权重衰减0.0001的AdamW优化器进行训练。在评估指标方面，作者使用了Synapse、ADCD和Polyp数据集的平均Dice相似系数（DSC）。此外，在Synapse数据集上还使用了平均95%的Hausdorff距离（HD）。HD指标提供了关于边界误差性能的更精确估计。DSC值范围从0到1，较大的值表示性能较好，而HD较小的值表示性能较好。

Synapse

作者在Synapse数据集上比较了GCtx-UNet与当前最先进的方法，包括U-Net、Att-UNet、TransUnet、SwinUnet、MISSFormer、TransDeepLab、HiFormer、GPA-TUNet和CS-UNet。表2总结了DSC和HD的性能比较。

在MedNet上预训练的GCtx-UNet具有第二高的平均DSC（82.39%）和第三低的平均HD（15.94毫米）。当在ImageNet上预训练时，GCtx-UNet在平均DSC（81.95%）和平均HD（16.8毫米）方面排名第四。请注意，这些指标与最佳的CS-UNet和HiFormer的平均DSC（83.27%）和平均HD（14.7毫米）相差不远，后者具有更高的计算复杂性。还要注意的是，GCtx-UNet显著优于基于CNN的方法UNet和Att-UNet。例如，UNet的平均DSC为76.85%，平均HD为39.70毫米，而Att-UNet的平均DSC为77.77%，平均HD为36.02毫米。GCtx-UNet的性能优于如Transdeeplab、MISSFormer和Swin-UNet等基于Transformer的方法，这些方法的平均DSC范围从79.13到81.96%，平均HD范围从18.20到21.25毫米。GCtx-UNet的性能与如TransUnet、GPA-TUNet、HiFormer和CS-UNet等混合模型的性能也具有竞争力，这些模型的平均DSC范围从77.48到83.27%，平均HD范围从14.70到31.69毫米。

图5显示了两个Synapse图像上的定性比较，其中GCtx-UNet的性能优于Swin-UNet和CS-UNet。GCtx-UNet正确分割了大部分器官，胆囊区域有少量误分类。相比之下，Swin-UNet过度分割了脾脏（一些属于脾脏的区域被误分类为左肾），而CS-UNet过度分割了胰腺。作者推测这种改进是由于使用了GC-ViT，它引入了一个参数高效的降采样模块以及修改后的融合MB-Conv块。这些修改解决了ViTs中的归纳偏差不足的问题，使GCtx-UNet能够准确捕捉相对较大的区域，并处理彼此靠近的器官。

对于具有清晰边界的大型器官（如肾脏、胰腺和脾脏）的分割需要网络捕捉全局特征。作者推测这就是为什么基于Transformer的模型比基于CNN的模型更准确的原因。对于像动脉这样的小型器官的分割则更多受益于局部特征的检测。这可能就是为什么基于CNN的模型比基于Transformer的模型结果更准确的原因。对于具有复杂边界的大型器官（如肝脏和胃）的分割需要捕捉局部和全局特征。这可能就是为什么混合模型结果更准确的原因。

Acdc

作者比较了GCtx-UNet在ACDC数据集上与一些最先进的方法的性能，包括基于CNN的方法（R50-UNet和R50-Atten-UNet）、基于Transformer的方法（R50-ViT、Swin-UNet、MISSFormer）以及混合方法（TransUNet、GPA-TUNet和CS-UNet）。结果如表3所示。在平均DSC方面，GCtx-UNet优于所有其他方法，其中在MedNet上预训练的GCtx-UNet性能最佳。特别是，在右心室（RV）和左心室（LV）方面，GCtx-UNet表现更佳。在MedNet上预训练的GCtx-UNet也比在ImageNet上预训练的GCtx-UNet表现更好。图6包含了ACDC数据集中的3张示例图像，用于定性比较GCtx-UNet和CS-UNet。在MetNet上预训练的GCtx-UNet能够比CS-UNet和在ImageNet上预训练的GCtx-UNet更准确地分割右心室和左心室。

Polyp datasets

作者首先在两个已知数据集（CVC-ClinicDB和Kvasir）上训练GCtx-UNet，然后在三个未见数据集（CVC-ColonDB、ETIS-LaribDB和CVC-300）上使用训练好的模型来评估模型的泛化能力。表4将GCtx-UNet的性能与基于卷积神经网络的最先进算法（UNet、UNet++、PraNet）和混合方法（CS-UNet）进行了比较。尽管在Kvasir数据集上，预先在MedNet上训练的GCtx-UNet具有最佳的DSC指标，但在CVC-ClinicDB数据集上它落后于CS-UNet和PraNet，不过差距相对较小。

泛化能力如表4（第3-5列）所示，GCtx-UNet对未见数据集（CVC-ColonDB、ETIS-LaribDB和CVC-300）具有更好的泛化能力。与其他方法相比，GCtx-UNet模型在CVC-ColonDB和ETIS-LaribDB数据集上取得了最佳和次佳的DSC值，在CVC-300数据集上取得了次佳的DSC值。GCtx-UNet的整体性能相当显著。图7展示了各种方法的定性分割结果，包括在MedNet和ImageNet上训练的GOctx-UNet。从每个数据集中选取了五个样本，以突出模糊边界和小息肉，便于区分分割性能。在MedNet上预先训练的GOctx-UNet显著减少了假阳性和假阴性。这种改进归功于它在区分息肉区域与正常区域的模糊边界方面的增强能力。

Conclusion

作者引入了Gctx-UNet，一种U形网络，它结合了轻量级的视觉 Transformer ，通过有效地捕捉全局和局部特征来增强医学图像分割。编码器和解码器组件之间的下采样和上采样块有效地提供了归纳偏置并建模通道间的依赖关系。

在多个医学图像数据集上，Gctx-UNet与传统基于CNN的方法、基于Transformer的方法以及混合方法相比，具有更好的或可比较的性能。同时，Gctx-UNet具有更低的模型复杂度，包括更少的模型参数数量、更低的模型大小、更低的训练和推理时间以及更低的训练FLOPs。

Gctx-UNet建模长距离空间依赖关系的能力以及其在分割复杂和小型解剖结构方面的竞争力，使其成为临床应用的有前景的工具。该架构的设计，包括带有跳跃连接的GC-ViT编码器和解码器，使其在保持低于现有技术水平计算复杂性的同时，还能提供高性能。

在医学图像数据集-MedNet上的预训练以及随后在多个医学成像任务上的评估显示了模型的鲁棒性和泛化能力，使Gctx-UNet成为一种实用且强大的医学图像分割方法。作为未来的工作，作者计划引入Gctx-UNet 3D模型用于医学图像的 Voxel 分割。

参考

[1].GCtx-UNet: Efficient Network for Medical Image Segmentation.