Transformer又一城 | Swin-Unet:首个纯Transformer的医学图像分割模型解读

简介: Transformer又一城 | Swin-Unet:首个纯Transformer的医学图像分割模型解读

1简介


在过去的几年中,卷积神经网络(CNN)在医学图像分析中取得了里程碑式的进展。尤其是,基于U形结构和skip-connections的深度神经网络已广泛应用于各种医学图像任务中。但是,尽管CNN取得了出色的性能,但是由于卷积操作的局限性,它无法很好地学习全局和远程语义信息交互。

在本文中,作者提出了Swin-Unet,它是用于医学图像分割的类似Unet的纯Transformer模型。标记化的图像块通过跳跃连接被送到基于Transformer的U形Encoder-Decoder架构中,以进行局部和全局语义特征学习。

具体来说,使用带有偏移窗口的分层Swin Transformer作为编码器来提取上下文特征。并设计了一个symmetric Swin Transformer-based decoder with patch expanding layer来执行上采样操作,以恢复特征图的空间分辨率。在对输入和输出进行4倍的下采样和上采样的情况下,对多器官和心脏分割任务进行的实验表明,基于纯Transformer的U-shaped Encoder-Decoder优于那些全卷积或者Transformer和卷积的组合。


2Swin-Unet架构


图1 Swin-Unet架构:由Encoder, Bottleneck, Decoder和Skip Connections组成。

Encoder, Bottleneck以及Decoder都是基于Swin-Transformer block构造的实现。

2.1 Swin Transformer block

图2 Swin Transformer block

与传统的multi-head self attention(MSA)模块不同,Swin Transformer是基于平移窗口构造的。在图2中,给出了2个连续的Swin Transformer Block。每个Swin Transformer由LayerNorm(LN)层、multi-head self attention、residual connection和2个具有GELU的MLP组成。在2个连续的Transformer模块中分别采用了windowbased multi-head self attention(W-MSA)模块和 shifted window-based multi-head self attention (SW-MSA)模块。基于这种窗口划分机制的连续Swin Transformer Block可表示为:

image.png

其中,和分别表示(S)W-MSA模块和第块的MLP模块的输出。

与前面的研究ViT类似,self attention的计算方法如下:

image.png

其中,表示query、key和value矩阵。和分别表示窗口中patch的数量和query或key的维度。value来自偏置矩阵。


2.2 Encoder

在Encoder中,将分辨率为的维tokenized inputs输入到连续的2个Swin Transformer块中进行表示学习,特征维度和分辨率保持不变。同时,patch merge layer会减少Token的数量(2×downsampling),将特征维数增加到2×原始维数。此过程将在Encoder中重复3次。

image.png

Patch merging layer

输入patch分为4部分,通过Patch merging layer连接在一起。这样的处理会使特征分辨率下降2倍。并且,由于拼接操作的结果是特征维数增加了4倍,因此在拼接的特征上加一个线性层,将特征维数统一为原始维数的2倍。


2.3 Decoder

与Encoder相对应的是基于Swin Transformer block的Symmetric Decoder。为此,与编码器中使用的patch merge层不同,我们在解码器中使用patch expand层对提取的深度特征进行上采样。patch expansion layer将相邻维度的特征图重塑为更高分辨率的特征图(2×上采样),并相应地将特征维数减半。

image.png

Patch expanding layer

以第1个Patch expanding layer为例,在上采样之前,对输入特征加一个线性层,将特征维数增加到原始维数的2倍。然后,利用rearrange operation将输入特征的分辨率扩大到输入分辨率的2倍,将特征维数降低到输入维数的1/4,即。


Up-Sampling会带来什么影响?

针对Encoder中的patch merge层,作者在Decoder中专门设计了Patch expanding layer,用于上采样和特征维数增加。为了探索所提出Patch expanding layer的有效性,作者在Synapse数据集上进行了双线性插值、转置卷积和Patch expanding layer的Swin-Unet实验。实验结果表明,本文提出的Swin-Unet结合Patch expanding layer可以获得更好的分割精度。

image.png

2.4 Bottleneck

由于Transformer太深导致收敛比较困难,因此使用2个连续Swin Transformer blocks来构造Bottleneck以学习深度特征表示。在Bottleneck处,特征维度和分辨率保持不变。

2.5 Skip connection

与U-Net类似,Skip connection用于融合来自Encoder的多尺度特征与上采样特征。这里将浅层特征和深层特征连接在一起,以减少降采样带来的空间信息损失。然后是一个线性层,连接特征尺寸保持与上采样特征的尺寸相同。


skip connections数量的影响?

Swin-UNet在1/4,1/8和1/16的降采样尺度上添加了skip connections。通过将skip connections数分别更改为0、1、2和3,实验了不同skip connections数量对模型分割性能的影响。从下表中可以看出,模型的性能随着skip connections数的增加而提高。因此,为了使模型更加鲁棒,本工作中设置skip connections数为3。

image.png


3实验结果


3.1 Synapse数据集

3.2 ACDC数据集

就很多基于Transformer的文章和已经开源的代码来看,模型的预训练会严重影响基于Transformer模型的性能。其实作者在文章的最后也提到了他们也是直接使用Swin-Transformer在ImageNet上的预训练权值来初始化网络的Encoder和Decoder的,这可能是一个次优方案。不过作者也在努力尝试实现完全端到端的实验了(此处点赞)。


4推荐阅读



5参考


[1].Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

[2].https://github.com/HuCaoFighting/Swin-Unet

相关文章
|
机器学习/深度学习 人工智能 数据管理
文生图的基石CLIP模型的发展综述
CLIP(Contrastive Language-Image Pre-training)是OpenAI在2021年发布的多模态模型,用于学习文本-图像对的匹配。模型由文本和图像编码器组成,通过对比学习使匹配的输入对在向量空间中靠近,非匹配对远离。预训练后,CLIP被广泛应用于各种任务,如零样本分类和语义搜索。后续研究包括ALIGN、K-LITE、OpenCLIP、MetaCLIP和DFN,它们分别在数据规模、知识增强、性能缩放和数据过滤等方面进行了改进和扩展,促进了多模态AI的发展。
2845 0
|
机器学习/深度学习 数据可视化 测试技术
YOLO11实战:新颖的多尺度卷积注意力(MSCA)加在网络不同位置的涨点情况 | 创新点如何在自己数据集上高效涨点,解决不涨点掉点等问题
本文探讨了创新点在自定义数据集上表现不稳定的问题,分析了不同数据集和网络位置对创新效果的影响。通过在YOLO11的不同位置引入MSCAAttention模块,展示了三种不同的改进方案及其效果。实验结果显示,改进方案在mAP50指标上分别提升了至0.788、0.792和0.775。建议多尝试不同配置,找到最适合特定数据集的解决方案。
3929 0
|
存储 数据采集 机器学习/深度学习
LIDC-IDRI肺结节数据集分割策略
本文介绍了使用LIDC-IDRI开源数据集进行肺癌检测项目的完整流程,包括数据预处理、训练分割模型和分类模型三个主要步骤。首先,下载包含患者DICOM文件的数据集;其次,克隆预处理代码并配置Pylidc库以生成肺部遮罩图像;最后,通过脚本准备数据集并创建元数据文件。文章还提供了相关GitHub资源链接,帮助读者更好地理解和实现项目。
1674 11
LIDC-IDRI肺结节数据集分割策略
|
机器学习/深度学习 计算机视觉
YOLOv11改进策略【YOLO和Mamba】| 2024 VM-UNet,高效的特征提取模块VSS block 二次创新提高精度
YOLOv11改进策略【YOLO和Mamba】| 2024 VM-UNet,高效的特征提取模块VSS block 二次创新提高精度
1484 9
|
机器学习/深度学习 数据可视化 自动驾驶
YOLO11-seg分割:具有切片操作的SimAM注意力,魔改SimAM助力分割
本文创新地对SimAM注意力机制进行魔改,引入切片操作,显著提升了小目标特征提取能力。针对SimAM在计算整张特征图的像素差平均值时可能忽略小目标重要性的问题,通过切片操作增强了小目标的加权效果。实验结果显示,魔改后的SimAM在YOLO11-seg上的Mask mAP50从0.673提升至0.681,有效改善了小目标检测性能。
1618 2
|
并行计算 数据挖掘 PyTorch
【YOLOv8改进 - 特征融合】DySample :超轻量级且高效的动态上采样器
【YOLOv8改进 - 特征融合】DySample :超轻量级且高效的动态上采样器
【YOLOv8改进 - 特征融合】DySample :超轻量级且高效的动态上采样器
|
机器学习/深度学习 人工智能 PyTorch
【人工智能】Transformers之Pipeline(七):图像分割(image-segmentation)
【人工智能】Transformers之Pipeline(七):图像分割(image-segmentation)
917 0
|
机器学习/深度学习 测试技术 网络架构
【YOLOv10改进-注意力机制】MSCAAttention多尺度卷积注意力
YOLOv10专栏介绍了一种新的卷积网络架构SegNeXt,它在语义分割任务中展现出优于Transformer模型的效率和效果。SegNeXt通过名为Multi-Scale Convolutional Attention (MSCA)的组件,结合局部信息聚合、多尺度上下文捕获和通道关系模拟,提升了性能。在多个数据集上,SegNeXt以较少参数实现了超过现有SOTA的性能,特别是在Pascal VOC 2012上,以1/10的参数量达到90.6%的mIoU。YOLOv10引入了MSCA模块,用于增强目标检测的上下文关注。相关代码和配置详情可在链接中找到。
|
机器学习/深度学习 资源调度 Java
【YOLOv8改进 - 注意力机制】GC Block: 全局上下文块,高效捕获特征图中的全局依赖关系
YOLOv8专栏探讨了目标检测的创新改进,如整合NLNet和SENet优势的GCBlock,用于高效全局上下文建模。GCNet在多个识别任务中表现优越,同时降低了计算成本。文章提供了论文、代码链接及详细实现,包括特征的全局建模、变换和融合步骤。核心GCBlock代码展示了其结构。更多实战案例和配置见相关链接。
|
机器学习/深度学习 编解码 自然语言处理
一文综述 | 万字文总结,近几年都有哪些语义分割模型用了Transformer方法呢?
一文综述 | 万字文总结,近几年都有哪些语义分割模型用了Transformer方法呢?
1627 0

热门文章

最新文章

下一篇
开通oss服务