原文地址：https://arxiv.org/pdf/2105.05537.pdf

作者：慕尼黑大学+复旦大学+华为。作者在华为实习所写

代码: 

GitHub - HuCaoFighting/Swin-Unet: The codes for the work "Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation"

首篇将transformers用于医疗分割的文章

设计的Idea：UNet + Swin Transformer(CVPR2021 best paper) ，使用的具体模块:SwinT+skip connection

Eecoder: pure transformer, not using CNN

Decoder: symmetrical to encoder

目录

摘要

1 Introduction

2 related work

2.1 CNN-based methods

2.2 Vision Transformers

2.3 Self-attention/Transformer to complement CNNs 

3 method 

3.1 Architecture overview

3.2 Swin Transformer block

3.3 Encoder

3.4 Bottleneck

3.5 Decoder

3.6 Skip connection

4 Experiments and Discussion

4.1 Datasets

4.2 Implementation details

4.3 Experiment results on Synapse dataset

4.4 Experiment results on ACDC dataset

4.5 Ablation study

Effect of up-sampling:

Effect of the number of skip connections:

Effect of input size:

Effect of model scale:

4.6 Discussion

参考推荐

摘要

我们提出了Swin-Unet，这是一种类似Unet的pure transformer，用于医学图像分割。tokenized image patches被送到基于 Transformer 的 带有skip-connections的U-型 Encoder-Decoder 架构中，用于局部与全局语义特征学习。具体来说，我们使用带有移位窗口shifted windows的分层 Swin Transformer 作为编码器来提取上下文特征。基于对称 Swin Transformer 的解码器具有patch expanding层，旨在执行上采样操作，以恢复特征图的空间分辨率。通过直接将输入输出进行4×的上采样和下采样，进行了多器官和心脏分割任务的实验。

1 Introduction

ViT、DeiT和Swin Transformer在图像 识别任务中的成功证明了Transformer在视觉领域的应用潜力。vision transformer（ViT）执行图像识别任务。ViT 以位置编码的 2D 图像patch作为输入，并在大型数据集上进行预训练，实现了与基于 CNN 的方法相当的性能。此外，data-efficient image transformer (DeiT) 表明Transformer可以在中型数据集上进行训练，并且可以通过将其与e distillation 相结合来获得更强大的Transformer。在[19]中，开发了一个分层的Swin Transformer， 在图像分类、对象检测和语义分割方面取得了最先进的表现（CVPR2021最佳论文，所以说本文站在了巨人的肩膀上）。

在Swin Transformer[19]成功的推动下，我们提出Swin-Unet，利用Transformer的强大功能进行2D医学图像分割。

2 related work

2.1 CNN-based methods

2.2 Vision Transformers

2.3 Self-attention/Transformer to complement CNNs 

近年来，研究人员试图将自我注意机制引入CNN，以提高网络的性能[13]。···但是，这仍然是基于CNN的方法。目前，CNN和Transformer的结合正在做出一些努力，以打破CNN在医学图像分割中的主导地位[2，27，1]。在[2]Trans-unet中，作者将Transformer与CNN相结合，构成了用于2D医学图像分割的强大编码器。····。与上述方法不同，我们尝试探索pure Transformer在医学图像分割中的应用潜力。就是说本文不与CNN做结合，直接将transformer作为编码器。

3 method 

3.1 Architecture overview

Swin-Unet的整体架构如图所示，由encoder, bottleneck, decoder 和3个 skip connections组成。Swin-Unet的基本单元是Swin Transformer block [19]。对于encoder，为了将输入转换为sequence embeddings，医学图像被拆分为非重叠的patches ，patch大小为 4 × 4。通过这种拆分方法，每个patch的feature dimension 变为 4 × 4 × 3 (通道数) = 48。因此从W，H，3的输入就转换为。此外，应用linear embedding层将输入投影到任意维度（表示为C)。转换后 patch tokens 通过多个 Swin Transformer 块和patch merging 层来生成hierarchical feature representations。具体而言，patch merging layer负责下采样和增加通道数，Swin Transformer块负责特征表示学习。

受 U-Net [3] 的启发，我们设计了transformer-based decoder。解码器由Swin Transformer块和patch expanding layer组成。提取的context features通过skip connections与编码器中的多尺度特征融合，以补充由下采样引起的空间信息丢失。与patch merging layer相比， patch expanding layer 执行上采样。将adjacent 维度的特征图reshape  2×上采样分辨率的更大特征图。最后一个patch expanding 用于执行 4×上采样，恢复到输入分辨率 （W ×H），然后在这些上采样特征上应用 a linear projection layer以输出像素级分割预测。我们将在以下部分中详细说明每个block。

3.2 Swin Transformer block

这个板块没有任何创新，就是直接套用SW-MSA和W-MSA。

不同与传统MSA模块，Swin Transformer block基于shifted windows。如图2，两个连续的Swin Transformer blocks被展示。每个block由LayerNorm (LN) layer, MSA module, residual connection 和 2-layer MLP with GELU non-linearity构成。window based multi-head self attention (W-MSA) module and the shifted window-based multi-head self attention (SW-MSA) module 分别应用于两个连续的transformer 模块。

基于这种窗口分割机制， 连续swin transformer 块可公式化为：

  一个接一个的套娃。

与前作[31，32]类似，SA的计算如下： 

3.3 Encoder

Swin Transformer blocks 不会改变特征大小。patch merging layer 下采样2倍tokens数量，增加2倍特征维度（C）.然后重复三次.

• Patch merging layer。将输入的patch分为4块，然后将他们catenated在一起。这样特征分辨率将变为1/2。并且特征维度（c）变成4×。最后加上一个linear layer，使得最后的维度减半，即变为原图的2倍。

3.4 Bottleneck

由于Transformer 太深会无法收敛，仅由2个 successive Swin Transformer blocks构成bottleneck。

3.5 Decoder

基于Swin Transformer模块，构建与encoder对称的Decoder。

3.6 Skip connection

与Unet相似，concatenate the shallow features and the deep features。接上一个a linear layer, the dimension of the concatenated features保持与the dimension of the upsampled features一样.

4 Experiments and Discussion

4.1 Datasets

Synapse multi-organ segmentation dataset (Synapse):

30 cases with 3779 axial abdominal clinical CT images. 18 samples for training set and 12 samples for validation set. evaluation metric: the average Dice-Similarity coefficient (DSC) and average Hausdorff Distance (HD) 

 pre trained parameter: Swin-T 

4.2 Implementation details

对于所有训练数据，使用翻转和旋转等数据增强来增加数据多样性。输入图片大小和 patch size 设为 224×224 and 4. 模型参数在ImageNet上预训练了。batch size= 24。popular SGD optimizer with momentum 0.9 and weight decay 1e-4。

4.3 Experiment results on Synapse dataset

比transUnet 好了不少。估计是swimTrans 帮了大忙。 此外多加了红线的三个模型作为对比。 

红框里的结果，swimUnet 表现好。

4.4 Experiment results on ACDC dataset

4.5 Ablation study

Effect of up-sampling:

测试了其他两种上采样方式的效果。用卷积的方式不如patch expand。

Effect of the number of skip connections:

0, 1, 2 and 3 分别对应1/4, 1/8, and 1/16 分辨率。可以看出1/4的提高有4个点，非常需要low level feature的融合、 

Effect of input size:

输入大小224 × 224 到 384 × 384，性能会有些提升，但是计算量大幅上升，所以选择224

Effect of model scale:

本文选择tiny one，因为计算量 

可以看到tiny的layer数量为 2 2 6 2

4.6 Discussion

两个改进点：

1、在这项工作中，我们直接使用 ImageNet 上 Swin transformer [19] 的训练权重来初始化网络编码器和解码器，这可能是一个次优方案。这种初始化方法很简单，将来我们将探索端到端预训练Transformer的方法，以进行医学图像分割。

2、由于本文中的输入图像是2D的，而大多数医学图像数据是3D的，因此我们将在未来探讨Swin-Unet在3D医学图像分割中的应用。

参考推荐：

代码：图像分割之Swin-Unet分享_xuzz_498100208的博客-CSDN博客

图解Swin Transformer - 知乎