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

打印 上一主题 下一主题

主题 909|帖子 909|积分 2727

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表现优秀,但由于卷积操作的局部性,其难以很好地学习全局和长距离语义信息的交互。在本文中,我们提出了Swin-Unet,一种用于医学图像分割的类Unet纯Transformer模型。被切分为Token的图像块输入到基于Transformer的U形编码器-解码器架构中,通过跳跃连接实现局部和全局语义特性的学习。具体来说,我们采取了带有移动窗口的分层Swin Transformer作为编码器以提取上下文特性;同时计划了一个基于对称Swin Transformer的解码器,结合Patch扩展层实现上采样操作,以规复特性图的空间分辨率。在输入和输出直接下采样和上采样4倍的环境下,针对多器官和心脏分割任务的实验表明,这种纯基于Transformer的U形编码器-解码器网络优于完全卷积方法或Transformer与卷积结合的方法。
二、模型

1、架构概述


Swin-Unet的团体架构如图所示。Swin-Unet由编码器瓶颈层解码器跳跃连接组成。Swin-Unet的基本单元是Swin Transformer块。
编码器
为了将输入转换为序列嵌入(将二维图像问题转换为序列问题),医学图像被分割为大小为
的不重叠块。通过这种分割方式,每个Patch的特性维度变为

接着,一个线性嵌入层被应用于将特性维度投射到任意维度(表示为
)。
转换后的Patch Token通过多个Swin Transformer块和Patch合并层,生成分层的特性表示。其中:


  • Patch合并层负责下采样和特性维度扩展,实现更高效的多层次特性提取。
  • Swin Transformer块负责特性表示的学习,捕获局部和全局上下文。
瓶颈层
连接编码器与解码器,进一步处理从编码器提取的深层次特性。
解码器
受U-Net的启发,计划了对称的基于Transformer的解码器。解码器由Swin Transformer块和Patch扩展层组成。解码器中提取的上下文特性通过跳跃连接与编码器的多标准特性融合,以增补因下采样导致的空间信息丢失。
与Patch合并层相对,Patch扩展层专为上采样计划:


  • 它将相邻维度的特性图重新分列为更大的特性图,实现分辨率的
    上采样。
  • 在末了一个Patch扩展层中,执行
    上采样以将特性图分辨率规复到输入图像的分辨率 (
    )。
  • 末了,对这些上采样后的特性应用一个线性投影层,以输出像素级的分割猜测。
2、Swin Transformer块

与传统的多头自注意力(MSA)模块不同,Swin Transformer块基于移动窗口(shifted windows)构建。如图所示,展示了两个连续的Swin Transformer块。每个Swin Transformer块由以下组件组成:

   (1)LayerNorm(LN)层:是一种归一化技术,通过对每一层的特性举行标准化(零均值和单元方差),加速模型训练,进步模型的稳固性。
  (2)多头自注意力模块(W-MSA):自注意力机制通过计算输入特性中每对元素之间的相关性(注意力权重),实现特性的动态加权组合,捕获全局依赖关系。多头机制(Multi-Head Attention)是自注意力的扩展版本,通过并行多个注意力头,更好地捕获不同子空间中的特性。
  (3)残差连接(Residual Connection):通过直接跳过非线性变换,将输入特性与输出特性相加,解决了深层网络中梯度消失的问题。
  (4)多层感知机(MLP):是一种简单的前馈神经网络,用于特性转换和非线性映射。在Swin Transformer块中,MLP包含两层全连接层,中央插入GELU激活函数。
  总结:
  (1)LayerNorm 提供了稳固的特性归一化。
  (2)多头自注意力模块 捕获全局上下文。
  (3)残差连接 保持梯度稳固,克制信息丢失。
  (4)带有GELU的两层MLP 增强了非线性表达能力,增补了注意力机制后的特性变换。
  在两个连续的Transformer块中,分别应用了基于窗口的多头自注意力(W-MSA)模块和基于移动窗口的多头自注意力(SW-MSA)模块。基于这样的窗口划分机制,连续的Swin Transformer块可以被公式化为:




其中,
分别表示第
个块中 (S)W-MSA 模块和 MLP 模块的输出。
自注意力机制的计算公式如下:

这里
 分别表示查询矩阵、键矩阵和值矩阵。
表示窗口中的Patch数目,
是查询或键的维度。矩阵
 的值取自偏置矩阵

3、编码器

在编码器中,分割为 
 维Token且分辨率为
的输入数据被输入到两个连续的Swin Transformer块中举行表示学习。在这一过程中,特性维度和分辨率保持稳固。同时,Patch合并层会将Token数目减少一半(即2倍下采样),并将特性维度增长到原始维度的2倍。编码器中会重复三次这一过程。
Patch合并层


  • 输入的特性被划分为4个子部门,这些子部门通过连接操作整合为一个新的特性表示。
  • 连接操作将特性的分辨率缩小了一半(2倍下采样),减少了空间上的计算复杂度。
  • 原始连接后的特性维度增长了4倍。为了保持特性表示的同等性,后续通过一个线性层将维度调整为原始特性维度的2倍。
4、瓶颈层

由于Transformer网络过深时难以收敛 ,因此瓶颈层仅利用了两个连续的Swin Transformer块构建,用于学习深层特性表示。在瓶颈层中,特性维度和分辨率保持稳固。
作为编码器息争码器之间的连接部门,瓶颈层提取和处理全局特性表示。通过减少网络深度(仅两个Swin Transformer块),克制过深Transformer模型带来的收敛问题,同时保持特性的高质量表达。
5、解码器

与编码器相对应,对称的解码器基于Swin Transformer块构建。为此,与编码器中利用的Patch合并层相对,解码器中利用Patch扩展层对提取的深层特性举行上采样。Patch扩展层将相邻维度的特性图重新分列为具有更高分辨率的特性图(2倍上采样),并将特性维度减少到原始维度的一半。
Patch扩展层
在上采样之前,对输入特性(
)应用一个线性层,将特性维度扩展到原始维度的两倍(
)。
利用重排(rearrange)操作将输入特性的分辨率扩展为输入分辨率的两倍,同时将特性维度减少到原始维度的四分之一:
从 
扩展到

解码器通过渐渐上采样规复特性分辨率,同时保持或减少特性维度,为最终的分割结果提供高分辨率上下文信息。
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个测试样本。仅利用均匀DSC作为评估指标来评估我们的方法在该数据集上的表现。
   Dice相似系数(DSC):衡量分割结果与真实标签之间的重叠水平,范围为0到1,值越大表示分割性能越好。
  

  其中, A 和B 分别是猜测分割结果和真实标签。
  Hausdorff距离(HD):衡量分割结果界限与真实界限之间的最大距离,用于评估界限细节的精确性。
  

   2、实现细节

Swin-Unet 基于 Python 3.6 和 Pytorch 1.7.0 实现。在所有训练案例中,为了增长数据多样性,利用了数据增强技术,如翻转和旋转。输入图像大小和Patch大小分别设置为 224×224和 4。我们在具有32GB内存的Nvidia V100 GPU上训练模型。模型参数利用在 ImageNet 上预训练的权重举行初始化。
在训练过程中,批量大小设置为24,优化模型利用了具有动量(momentum)0.9和权重衰减(weight decay)
的流行SGD优化器来执行反向传播。
3、 Synapse数据集上的实验结果


表中展示了所提出的Swin-Unet与以往的最新方法在Synapse多器官CT数据集上的对比结果。与TransUnet 不同,我们添加了自己实现的U-Net 和 Att-UNet在Synapse数据集上的测试结果。实验结果表明,我们基于纯Transformer的类Unet方法取得了最优的分割性能,达到了79.13%的分割正确率(DSC)和21.55%的Hausdorff距离(HD)。与Att-Unet  和最新的TransUnet方法相比,尽管我们的方法在DSC评价指标上没有明显提升,但在HD评价指标上分别取得了约4%和10%的精度提升,这表明我们的方法在界限猜测方面表现更优。

图中展示了不同方法在Synapse多器官CT数据集上的分割结果。从图中可以看出,基于CNN的方法容易出现过分割问题,这可能是由于卷积操作的局部性导致的。在本研究中,我们通过将Transformer集成到具有跳跃连接的U形架构中,证实了无卷积的纯Transformer方法可以更好地学习全局和长距离语义信息交互,从而实现更优的分割结果。
4、ACDC 数据集上的实验结果


与 Synapse 数据集类似,我们在 ACDC 数据集上训练了所提出的 Swin-Unet 以执行医学图像分割任务。实验结果总结在表中。利用 MRI 模态的图像数据作为输入,Swin-Unet 依然能够取得 90.00% 的良好分割正确率,这表明我们的方法具有良好的泛化能力和鲁棒性。
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×224的输入分辨率。

(4) 模型规模的影响
从表中可以看出,模型规模的增长几乎未能明显进步模型性能,但却增长了整个网络的计算本钱。
综合思量正确性和速度的平衡,我们选择基于 Tiny 的模型执行医学图像分割任务。


免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!更多信息从访问主页:qidao123.com:ToB企服之家,中国第一个企服评测及商务社交产业平台。

本帖子中包含更多资源

您需要 登录 才可以下载或查看,没有账号?立即注册

x
回复

使用道具 举报

0 个回复

倒序浏览

快速回复

您需要登录后才可以回帖 登录 or 立即注册

本版积分规则

尚未崩坏

金牌会员
这个人很懒什么都没写!

标签云

快速回复 返回顶部 返回列表