UNet家族最强系列 | UNet、UNet++、TransUNet与SWin-UNet究竟哪个更强！！！

点击下方卡片，关注「集智书童」公众号

本文首发于【集智书童】，白名单账号转载请自觉植入本公众号名片并注明来源，非白名单账号请先申请权限，违者必究。

医学图像分割是医学图像分析的重要步骤，特别是作为高效疾病诊断和治疗的关键先决条件。深度学习在图像分割中的应用已成为一种普遍趋势。目前广泛采用的方法是U-Net及其变种。此外，随着预训练模型在自然语言处理任务中的显著成功，基于Transformer的模型（如TransUNet）在多个医学图像分割数据集上取得了令人满意的性能。
在本文中，作者对近年来最具代表性的4种医学图像分割模型进行了调查。作者在理论上分析了这些模型的特点，并在两个基准数据集（结核病胸部X射线和卵巢肿瘤）上定量评估了它们的性能。最后，作者讨论了医学图像分割领域的主要挑战和未来趋势。作者的工作可以帮助相关领域的研究人员快速建立针对特定区域的医学分割模型。

1、简介

随着医学影像技术的不断发展，医学图像在疾病诊断和治疗计划中变得至关重要。医学图像分割在医学图像分析的基础和关键技术中起着重要作用。医学图像分割是指从CT或MRI等医学图像中识别器官或病变像素。它是医学图像分析中最具挑战性的任务之一，旨在传达和提取有关这些器官或组织的形状和体积的关键信息。

传统的医学图像分割方法主要依赖于医生进行手动特征提取，或基于图像处理技术和数学模型的手工设计，如阈值处理、边缘检测和形态学操作。这些方法提供了一定程度的可解释性和可控性。然而，由于医学图像的复杂性和多样性，以及医学图像分割任务的特异性，传统分割方法存在一定的局限性。手工算法在处理大量用于分割任务的医学图像时，无法满足效率和准确性的要求。此外，从医学图像中手动提取特征需要具有丰富专业知识和经验的医生，使他们容易受到主观因素的影响。

近年来，深度学习技术广泛应用于医学图像分割，以解决上述问题。通过深度特征学习，模型可以从图像中提取语义信息，从而提高分割的准确性，并灵活适应不同的医学图像数据集和任务。基于卷积神经网络（CNNs）的分割模型取得了显著的结果。例如，U-Net模型在ISBI 2015细胞分割挑战赛中获得第一名，SegNet模型在CamVid数据集上的语义分割任务中表现出色，等等。然而，卷积神经网络对于长距离依赖性的建模能力有限，使得充分利用图像内的语义信息变得具有挑战性。

最近，一些新的分割模型被提出，包括TransUNet和Swin-Unet。TransUNet是一种引入了Transformer模块的分割模型，以提高模型对长距离依赖性的建模能力。Transformer模块采用自注意力机制，计算输入序列中每个位置与其他位置之间的相似性，从而得到一个权重向量。这个权重向量用于计算每个位置的加权表示，促进全局信息的交互和整合。

换句话说，Transformer模型可以通过自注意力机制有效地捕捉输入序列中不同位置之间的关联性，从而更好地理解和处理序列数据。在TransUNet中，Transformer模块嵌入在U型架构内，从图像中提取全局信息，增强了模型的语义表示能力，并使其更适合处理大尺寸、高分辨率的医学图像。

另一方面，Swin-Unet是另一种引入了Swin Transformer模块的新型分割模型，以提高计算效率。Swin Transformer是一种分层自注意力机制，将输入特征图分解为多个Patches，每个Patch独立计算注意力权重，从而降低计算复杂度。Swin-Unet中的Swin Transformer模块与U型架构相结合，允许从图像中提取全局信息，同时降低计算复杂度和内存消耗。这使得它更适用于医学图像分割任务。

尽管近年来医学图像分割技术发展迅猛，但在深度学习模型在医学图像分割中的应用方面，仍然缺乏关于最新分割模型引入和这些模型之间定量性能比较的综合综述。

本文对近年来最具代表性的四种医学图像分割模型进行了调查：U-Net、UNet++、TransUNet和Swin-Unet。分析了这些模型的特点，并在两个基准数据集上进行了定量评估。最后，作者讨论了医学图像分割领域的主要挑战和未来发展趋势。此外，作者在GitHub上分享了所有实验源代码和详细的模型配置参数，以帮助相关研究人员快速了解这些模型，并对新的分割任务进行建模。

2、典型的医学图像分割模型

近年来，借助深度学习，医学图像分割取得了巨大的进展。卷积神经网络（CNNs），尤其是完全卷积网络（FCNs），在医学图像分割领域占据主导地位。随着医学图像分割的发展，在不同的模型变种中，U-Net已经成为事实上的选择，它由对称的编码器-解码器网络组成，并具有增强的跳跃连接，以增强细节保留。

基于这个神经网络，图像特征可以自动提取并用于分割任务。在医学图像分割中，已经使用了几种深度学习模型并取得了优异的结果，例如U-Net，UNet++，3D U-Net，V-Net，Attention-UNet，TransUNet和Swin-Unet。

2.1、U-Net

U-Net是医学图像分割模型中最知名的网络架构之一。它是由Ronneberger等人在2015年的ISBI挑战中提出的。U-Net模型被认为是医学图像分割中的经典模型，已广泛应用于各种任务，包括CT、MRI和X射线分割。模型结构如图1所示。其成功在于将卷积神经网络（CNNs）的深度特征提取能力与全卷积网络（FCNs）的像素级分割能力相结合。它还结合了跳跃连接等技术，以利用低层和高层特征信息，从而提高了分割的准确性和稳健性。

U-Net网络由收缩路径和扩展路径组成。收缩路径遵循典型的卷积网络架构。在每个下采样步骤中，特征通道的数量会加倍。扩展路径中的每个步骤包括上采样特征图，减半特征通道的数量，并将它们与来自收缩路径的相应裁剪特征图进行连接。在最后一层，使用1x1卷积将每个64分量的特征向量映射到所需的类别数。网络总共由23个卷积层组成。

自U-Net模型引入以来，基于UNet的几个改进版本已经出现，包括UNet++、Attention-UNet、TransUNet和Swin-Unet等。这些模型在原始U-Net模型的优势基础上进一步增强了分割性能，引入了注意力机制、转换网络结构和其他技术。因此，U-Net模型在医学图像分割中占据着重要的地位和影响力。

2.2、UNet++

UNet++网络架构是由周等人于2018年提出的，将密集连接的概念引入了U-Net网络中。模型结构如图2所示。UNet++在保留长跳跃连接的基础上，增加了更多的短跳跃连接路径和上采样卷积块，形成了新的编码器层次。UNet++中的U形连接结构是通过将解码器中的每个编码器与相同层次的其他编码器融合实现的。

具体而言，每个编码器从其他编码器接收相同尺度的特征图，并将它们连接在一起，以获得更具辨别性的特征表示。此外，后来提出的Attention-UNet++通过在编码器融合过程中添加注意机制，改进了特征图的连接，以增强对重要特征的关注和提取。

UNet++通过引入密集连接从不同层次捕获特征，实现了从不同层次和尺度提取特征信息。这些特征被整合到最终的预测中，以提高分割的准确性。密集连接的思想源自DenseNet。在DenseNet之前，卷积神经网络的演进通常涉及增加网络的深度或宽度。

DenseNet通过重复使用特征引入了一种新的结构，不仅缓解了梯度消失的问题，还减少了模型参数的数量。在原始的U-Net网络架构中，使用中间隐藏层的深度监督解决了UNet++训练过程中的梯度消失问题。这还允许在测试阶段对网络进行剪枝，从而减少模型的推理时间。

2.3、TransUNet

TransUNet网络架构是由陈等人于2021年提出的，是一种基于Transformer的分割网络。模型结构如图3所示。TransUNet在U-Net模型的基础上引入了混合编码器，将CNN和Transformer结合起来，以解决传统卷积神经网络在建模长距离依赖性和处理大尺寸图像方面的局限性。TransUNet的核心是Transformer模块，它包括多头自注意力机制和前馈神经网络。多头自注意力机制捕获图像中不同位置之间的依赖关系，建立特征表示中的全局上下文信息。这使得TransUNet能够更好地处理长距离依赖性，捕获图像中的语义信息，并提高模型的表示能力和泛化性能。

具体而言，TransUNet首先使用CNN提取特征并生成输入图像的特征图。然后，将这些特征图划分为大小为1x1的块，并馈送到另一个由12个Transformer模块组成的堆栈中。这种混合结构将卷积神经网络的特征提取能力与使用Transformer模块进行有效的全局信息建模相结合，相比于仅使用Transformer作为编码器，能够获得更好的性能。

TransUNet中的解码器对编码特征进行上采样，并将其与高分辨率的CNN特征图结合起来，以丰富语义信息，实现更精确的定位。最后一步涉及将特征图恢复到原始图像尺寸，并生成像素级的分割结果。与使用卷积神经网络的传统U形模型相比，TransUNet引入了一个由12个Transformer模块组成的堆栈，显著增加了参数数量，增加了模型训练的难度。在这项研究中，为了满足TransUNet在GPU上的训练需求，采用了降低批量大小的次优方法。

2.4、Swin-Unet

Swin-Unet网络架构是由曹等人于2023年提出的。模型结构如图4所示。与Trans-Unet不同，后者将U-Net编码器中的卷积块替换为Transformer块，Swin-Unet则利用Swin Transformer块来从输入图像中提取分层特征。Swin-Unet是第一个纯Transformer-based的U形架构。Swin Transformer将传统Transformer的一维序列扩展为二维图像块，并采用分层注意机制来捕获更大感受野范围内的特征。这个结构类似于卷积神经网络中的分层结构，用于特征提取。此外，Swin Transformer引入了在自注意力机制之上的窗口平移机制。通过将注意力计算限制在当前区域附近的窗口内，Swin-Unet更好地保留位置信息，进一步提高了模型的性能。

在Swin-Unet中，Swin Transformer被应用于编码、瓶颈和解码模块。值得注意的是，Swin-Unet中每个层的特征压缩比TransUNet要小。Swin-Unet并没有添加额外的Transformer模块，而是用Transformer模块替换了卷积模块，从而有效减少了模型参数的数量。

总体而言，Swin-Unet充分利用了Swin Transformer和U-Net的优势，为医学图像分割提供了一个有前途的方法。它在各种分割挑战和基准测试中展现出竞争性的性能。

3、实验

3.1、训练超参

3.2、实验结果

首先，作者在公开可用的数据集——结核病胸部X射线数据集上评估了每个模型的性能。表3呈现了每个模型的实验结果，而图7提供了模型性能的可视化表示。

结果表明，TransUNet模型在所有六个度量指标上取得了最佳表现，分别为96.45%（DSC↑）、10.75（HD↓）、93.25%（IoU↑）、98.16%（Acc↑）、97.36%（Precision↑）和95.72%（Recall↑）。此外，所有模型在肺部分割任务中均表现出色，对于所有四种分割方法，mIoU（平均交集联合）值均超过91%。在测试集中，剔除样本不清晰采样较少的情况下，剩余样本的DSC分数超过82%。根据它们的分割结果，所有四种方法都能有效满足分割要求。

接着，作者在具有不同成像模态的数据集上进一步评估了每个模型的性能。表4呈现了卵巢肿瘤数据集的结果。

作者根据平均Dice系数和平均IoU来评估实验结果。TransUNet模型表现最佳，其次是Swin-Unet、UNet++和U-Net。当基于平均Hausdorff距离评估结果时，TransUNet模型也表现最佳，其次是U-Net、Swin-Unet和UNet++。在平均准确率方面，TransUNet模型表现最佳，其次是Swin-Unet、U-Net和UNet++。基于平均精度评估结果时，TransUNet模型表现最佳，其次是U-Net、Swin-Unet和UNet++。

值得注意的是，Trans-UNet模型的精确度明显高于其他模型，表明其在准确识别肿块区域方面具有很强的能力。对于平均召回率，Swin-Unet模型表现最佳，其次是TransUNet、UNet++和U-Net。

图8提供了模型性能的可视化表示。综合考虑所有指标和实际的分割结果，Trans-UNet模型在所有五个评估指标上表现最佳，分别为89.18%（DSC↑）、22.35（HD↓）、82.73%（IoU↑）、99.02%（ACC↑）和92.28%（Recall↑）。其预测结果与实际标签在很大程度上相似。另一方面，Swin-Unet、UNet++和U-Net的评估指标较低，在实际预测中并未产生完美的结果。

作者进一步检查了每个模型的分割结果，当Dice系数极低（小于20%）时。作者将这些结果视为完全不可接受的分割结果。这些结果的统计数量显示在表5中。结果显示，在引入Transformer模块后，模型能够捕获全局信息，因此在卵巢肿块的主要区域上展示了更少的完全错误判断的情况。

3.3、讨论

在作者的实验中，所有模型都采用了U型架构，U-Net模型的引入在医学图像分割中具有重要意义。U-Net将编码器和解码器结合在一起，通过利用不同尺度的信息并保留高分辨率特征来实现精确分割。这种设计使得U-Net能够在医学图像分割任务中获得更准确的结果，提高病变的定位和分割精度。

此外，U-Net具有良好的可扩展性，可以通过添加或调整网络层、修改网络结构等来进行改进。这种灵活性使得U-Net模型能够应用于各种医学图像分割任务，并与其他深度学习模型集成和优化。

起初，Transformer模型在医学图像分割方面并不被认为具有潜力，因为它们天生缺乏定位能力。然而，TransUNet引入了一种结构，将Transformer与卷积神经网络相结合，形成了一个有效的编码器，并提高了分割性能。另一方面，如果不将卷积网络与Transformer结合使用，最终的结果不理想，就像Swin-Unet所示。

至于数据集，卵巢肿块的分割非常具有挑战性，图像的分割区域通常在大小、形状、位置和纹理上都存在差异，使得检测肿块更加困难。在肺部分割数据集中，识别肺部的范围相对容易，因为与背景相比，它具有明显的特征，每个模型在肺部分割任务中表现出色。

4、数据集挑战和问题

在医学图像处理中应用监督学习面临的最大挑战是医学图像标注。监督学习需要大量带有注释的样本输入，以获得良好的性能和稳定的泛化能力。然而，收集如此庞大的带注释病例数据集通常是一项非常艰巨的任务。医学图像需要专业临床医生的解释来收集、标注和注释。医学图像处理中的另一个重要问题是数据不平衡。在不平衡的数据集中，各类别之间的类别分布是不对称的；例如，在卵巢肿块数据集中，良性和恶性患者的数量存在自然的不平衡（异常患者的数量大于正常患者）。不平衡的数据可能会严重影响模型的性能。

已经广泛使用了一些方法来解决上述问题。数据增强可以通过对样本应用一组仿射变换，如翻转、旋转、镜像以及增强颜色（灰度）值，从而增加训练数据集的数量并平衡正负样本的比例。从在相同或其他领域中实现的成功模型进行迁移学习是上述问题的另一种解决方案。与数据增强相比，迁移学习是一种更具体的解决方案，只需要适度的计算资源和较少的标注数据，就能显著降低医学图像分割的错误率。

5、总结

在本文中，作者首先对医学分割进行了一般性介绍，然后研究了4种最具代表性的医学图像分割模型，即U-Net、UNet++、TransUNet和Swin-Unet。此外，作者在两个基准数据集上对这些模型进行了定量性能评估。为了帮助相关领域的研究人员快速理解这些模型并对新的分割任务进行建模，作者在GitHub上共享了所有的实验源代码和详细的模型设置参数。

近年来，基于机器学习的图像分割发展迅速。Meta 提出了SAM，彻底改变了图像分割。这是首次在图像分割中引入基础模型的概念，实现了零迁移的图像分割。与以往只能处理特定类别图像的分割模型不同，SAM可以处理所有图像，并通过提示实现准确的图像分割。Ma等人提出了MedSAM用于通用图像分割，这是首次将SAM应用于医学领域，并在医学分割任务上胜过了默认的SAM模型。基于大模型的图像分割已成为图像分割的未来趋势，这是一个具有广阔前景的有希望的研究方向。

在未来的工作中，作者将大模型的图像分割与医学图像相结合，通过零迁移的分割模型来解决难以获取的医学图像，从而为医学图像分析做出贡献。