《Swin-Unet: 类Unet的纯Transformer架构用于医学图像分割》学习笔记

paper：https://arxiv.org/pdf/2105.05537

GitHub：https://github.com/HuCaoFighting/Swin-Unet

目录

一、摘要

二、模型

1、架构概述

2、Swin Transformer块

3、编码器

4、瓶颈层

5、解码器

6、跳跃连接

三、实验

1、数据集

2、实现细节

3、 Synapse数据集上的实验结果

4、ACDC 数据集上的实验结果

5、消融实验

---

一、摘要

近年来，在医学图像分析领域中卷积神经网络（CNN）已显示出显著的进展。特别是在采用U型架构与跳跃连接的设计下，深度神经网络已被成功应用于多种医学图像处理任务中。然而尽管其表现出色，在CNN方法中由于其局部性质限制了其对全局语义特征与其他远处语义信息之间相互作用的有效学习能力。为了改进这一局限性，在本文中我们提出了一种新的纯Transformer架构Swin-Unet模型专门用于医学图像分割任务。该模型将输入图像划分为Token块后输入到基于Transformer架构设计的编码器-解码器框架中通过跳跃连接机制实现对局部与全局语义特征的有效学习过程在此过程中我们采用了带移动窗口机制的分层Swin Transformer作为编码器部分以提取多尺度上下文信息；同时设计了一个基于对称Swin Transformer构建的解码器模块结合Patch扩展层实现上采样操作从而恢复了特征图的空间分辨率特性值得注意的是在输入与输出端分别实施了4倍的空间采样率缩减与还原操作针对多器官分割及心脏分割等典型任务进行实验研究结果表明所提出的纯Transformer架构比完全依赖于CNN的方法或者混合使用Transformer与CNN的方法表现更为优异

二、模型

1、架构概述

该系统的核心结构可见于图中。该系统由编码器模块、瓶颈层组件、解码器部分以及跨越连接构成。其核心组件为基于旋转加注意力机制的模块。

编码器 ：

在将输入编码为序列嵌入的过程中（具体而言是将其转化为一个二维图像问题），医学图像被划分为不同尺寸的区域

的不重叠块。通过这种分割方式，每个Patch的特征维度变为

。

接着，一个线性嵌入层被应用于将特征维度投射到任意维度（表示为

）。

改写说明

• Patch合并层 主要通过下采样与扩展特征维度来提高数据的层次化特征提取能力。
• Swin Transformer块 主要承担着从局部到全局范围内的上下文学习任务。

瓶颈层 ：

连接编码器与解码器，进一步处理从编码器提取的深层次特征。

解码器 ：

参考了U-Net的设计理念后，在开发了一种对称式的基于Transformer架构的解码器的基础上构建了解码模块体系

与Patch合并层相对，Patch扩展层专为上采样设计：

• 它将相邻维度的特征图重新排列为更大的特征图，实现分辨率的

上采样。

• 在最后一个Patch扩展层中，执行

上采样以将特征图分辨率恢复到输入图像的分辨率 (

在此时将经过上采样的特征通过线性投影层进行处理，并最终实现像素级别的分割预测

2、Swin Transformer块

与传统的多头自注意力（MSA）模块不同的是, 被提出的Swin Transformer 块采用了移窗机制（shifted windows）进行结构化设计. 如图所示, 该文随后详细展示了两个连续排列的Swin Transformer 模块. 每个 Svin Transformer 块均由以下几个核心组件构成:

1. LayerNorm（LN）层：一种标准化技术，在每一层引入零均值和单位方差的标准化处理过程后，在促进模型训练效率的同时提升了模型稳定性。
2. 多头自注意力模块（W-MSA）：基于输入特征间关联程度计算出注意力权重后实现动态特征组合机制，在捕捉长距离依存关系方面具有显著优势。多头机制则通过并行运行多个独立的注意力头来增强不同子空间信息的捕捉能力。
3. 残差连接（Residual Connection）：通过省去非线性变换步骤直接将输入与输出拼接起来的方式有效缓解了深层网络中的梯度消失问题。
4. 多层感知机（MLP）：作为前馈神经网络的一种基本结构，在Swin Transformer模块中由两层全连接层构成并结合GELU激活函数实现复杂的非线性变换功能。
总结：
1. LayerNorm 提供了稳定的特征标准化方法。
2. 多头自注意力模块 在捕捉长距离依存关系方面表现出色。
3. 残差连接机制 保证了梯度传输的有效性和稳定性。
**4. 带有GELU激活函数的两层MLP架构显著增强了模型表达能力。

在两个连续的Transformer层中依次采用了基于滑动窗口的多头自注意力（W-MSA）模块和基于固定滑动窗口的多头自注意力（SW-MSA）模块。通过这种滑动窗口划分策略能够用数学公式表达为：

其中，

和

分别表示第

个块中 (S)W-MSA 模块和 MLP 模块的输出。

自注意力机制的计算公式如下：

这里

分别表示查询矩阵、键矩阵和值矩阵。

表示窗口中的Patch数量，

是查询或键的维度。矩阵

的值取自偏置矩阵

。

3、编码器

在编码器中，分割为

维Token且分辨率为

在输入数据经过两个连续的Swin Transformer块的过程中完成表示学习任务。在此过程中，特征维度与分辨率均保持恒定状态。随后，在Patch合并层中实现Token数量减半（即2倍下采样），同时将特征维度提升至原始维度的两倍。编码器模块经过三次重复迭代以完成整个特征提取流程。

Patch合并层 ：

• 输入将被划分为四个子部分，并通过连接操作将其整合成一个新的特征表示。
• 该连接操作使输入的空间分辨率减半（经过2倍下采样处理），从而降低了空间计算复杂度。
• 在经过连接操作后，原始维度增长了四倍。为了维持统一的尺度特性，在后续步骤中我们引入了一个线性层来将其调整回原始尺寸的一半。

4、瓶颈层

因为深度较大的Transformer网络在训练过程中难以达到稳定的收敛状态，研究者们选择将注意力机制模块进行优化设计。具体而言，在模型架构中采用了两个连续构建的Swin Transformer模块来增强深层特征表示能力。值得注意的是，在此设计框架下，并未影响整体模型的整体性能指标。

作为编码器与解码器之间的重要连接模块，在模型架构中负责提取并处理全局特征表示。通过缩减网络深度（仅引入两个Swin Transformer模块），有效规避过深Transformer架构所带来的收敛挑战，并持续保持高质量的特征表达能力。

5、解码器

对应于编码器的设计, 解码器部分采用了基于Swin Transformer块的对称架构。在实现过程中, 通过特定设计实现了特征信息的有效上采样。通过特定设计实现了相邻维度的空间信息重新排列, 将其转换为更高分辨率（2倍）的空间表示，并将特征维度减少到原始的一半

Patch扩展层 ：

在上采样之前，对输入特征（

）应用一个线性层，将特征维度扩展到原始维度的两倍（

）。

通过重排过程提升至加倍后的分辨率，并使特征维度降低至25%

从

扩展到

。

该解码模块通过分阶段放大操作逐步提升细节层次，并在不增加额外计算负担的情况下维持或降低输入特征的空间维度。

6、跳跃连接

与其U-Net相似，在此方案中采用跳跃连接来整合编码模块的多尺度特性和解码模块的上采样特性。

作用：

• 编码模块在经过多级降采样后可能会导致空间分辨率信息的缺失，并通过跳跃连接机制将回补至高分辨率特征。
• 通过融合编码器提取出的细节特征与解码器捕捉到的整体语义特征（即全局语义信息），从而提高分割精度。

三、实验

1、数据集

Synapse多器官分割数据集 (Synapse) ：

该数据集包含了30例病例，并总共拥有3779张轴向腹部临床CT图像。其中18个样本被分配给训练组而12个样本被分配给测试组。我们采用了包括平均Dice相似性指标（DSC）以及平均Hausdorff距离（HD）在内的两个评估指标，并对包括主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺以及胃在内的8个腹部器官进行了评估分析。

自动化心脏诊断挑战数据集 (ACDC) ：

ACDC 数据集来源于不同患者使用的 MRI 扫描设备。 在对每位受检者的 MR 影像进行解剖学标记时，在每个切片中分别标注了左心室 (LV)、右心室 (RV) 以及心肌层 (MYO)。 在构建该数据集时采用了 70% 的数据用于训练、10% 用于验证以及 20% 用于测试。 为了评估我们方法的性能，在分析该数据集时仅计算并报告平均 Dice 系数 (DSC) 。

Dice相似系数（DSC）是用来评估两个分割输出之间重叠程度的一种指标。其取值范围为0到1，在此范围内数值越大则表示该分割方法与标注信息的重叠程度越高。

其中，
A
和
B
分别代表预测分割结果和真实标签。
Hausdorff
距离（HD）
被用来评估分割结果边界与真实边界之间的最远距离，
它能够有效地反映区域细节的精确性。

2、实现细节

基于Python开发的Swin-Unet依赖于PyTorch框架1.7.0版本。针对所有训练样本集的数据多样性需求，在该系统中采用了图像翻转与旋转等数据增强技术。实验中设定输入图像尺寸与Patch划分尺寸分别为224×224像素及4px。该研究工作在配备32GB内存的NVIDIA Quadro V100显卡上进行了深度学习模型的构建与训练。网络参数初始值采用了从ImageNet大数据集上预训练得到的权重值。

在模型训练阶段，在采用批量大小为24的情况下进行优化操作。该优化过程采用了动量因子为0.9以及L2正则化系数来进行参数更新计算。

的流行SGD优化器来执行反向传播。

3、 Synapse数据集上的实验结果

表中列出了所提出的Swin-Unet与现有最新技术在Synapse多器官CT数据集上的对比实验结果。相较于TransUnet方法，在本研究中我们还引入了自研U-Net架构和Att-UNet模型，并在同一数据集上进行了测试评估。实验结果显示，通过纯Transformer架构设计的类U-Net模型展现出最佳分割性能，在DSC指标上达到79.13%，同时Hausdorff距离（HD）值为21.55%。尽管在DSC指标上的提升并不显著，在HD指标方面却分别实现了约4%和10%的精度提升。这表明我们的模型在边界预测任务中的表现更为出色。

图中呈现了不同方法在Synapse多器官CT数据集上的实验结果对比。观察到基于CNN的方法往往会出现明显的过分割现象，这一问题可能源于其卷积操作所具有的局限性。本研究采用将Transformer技术整合于带有跳跃连接的传统U形架构之中，并成功证明了完全去除卷积操作后仅依靠Transformer模型便能有效提取全局及远距离语义关联的能力，在此过程中实现了分割效果的重大提升。

4、ACDC 数据集上的实验结果

与Synapse数据集相仿，在ACDC数据集上基于所提出的Swin-Unet模型进行了相关实验以完成医学图像分割任务。实验结果已归纳整理于表中。通过使用MRI模态的数据输入到该模型中进行训练后发现，在测试集中达到了高达90%的精确度目标，并且进一步验证了该方法具备良好的泛化能力和稳定性

5、消融实验

为了探究不同因素对模型性能的作用，在 Synapse 数据集中进行了消融实验分析，并涉及了如下因素：上采样策略、跳跃连接规模、输入维度以及模型参数数量。

（1）上采样方式的影响 ：

为了适应编码器输出的Patch合并层，在解码器部分特别增加了Patch扩展层，并以实现图像放大并提升特征多样性为目标进行设计。

为了考察所提出的Patch扩展层的表现力，在Synapse数据集上分别进行了双线性插值、反卷积以及Patch扩展层的实验比较。

通过实验数据测试发现，在图像分割任务中将Patch增强模块集成到Swin-UNet架构中能够显著提升模型性能，并验证了该方法在实际应用中的有效性

（2）跳跃连接数量的影响 ：

Swin-Unet 的跳跃连接添加在分辨率为 1/4、1/8和 1/16的位置。

该研究依次将跳跃连接的数量配置为 0、1、2 和 3，并评估不同配置对模型分割效果的影响。

从表格数据可以看出，模型性能随着跳跃连接数量的增加而提高。由此可见，在本研究中将跳跃连接的数量设定为3个能够有效提升模型的鲁棒性。

（3） 输入大小的影响 ：

表展示了 Swin-Unet 使用 224×22和 384×384 输入分辨率时的测试结果。

当输入尺寸从 224×224 升高至 384×384时，在维持Patch尺寸为4的前提下，在Transformer架构下其输入Token序列长度随之增加，并显著提升了模型在分割任务中的性能表现。

然而, 尽管模型分割精度有所提升, 但这会导致整个网络计算负担明显加重. 为保证算法运行效率, 本研究中的实验均基于 224\times224 的输入分辨率设置.

（4） 模型规模的影响 ：

通过表格数据可以看出，基本模型规模的增加未能带来显著提升，但计算开销有所上升。

在确保准确性的同时兼顾处理速度的情况下，在医学图像分割任务中采用 Tiny 模型进行处理。