Transformer再下一城！Swin-Unet：首个纯Transformer的医学图像分割网络

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

本文转载自：集智书童

Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

论文：https://arxiv.org/abs/2105.05537

代码：https://github.com/HuCaoFighting/Swin-Unet

首个基于纯Transformer的U-Net形的医学图像分割网络，其中利用Swin Transformer构建encoder、bottleneck和decoder，表现SOTA！性能优于TransUnet、Att-UNet等，代码即将开源！
作者单位：慕尼黑工业大学, 复旦大学, 华为(田奇等人)

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可表示为:

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

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

其中， 表示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次。

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×上采样)，并相应地将特征维数减半。

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可以获得更好的分割精度。

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。

3实验结果

3.1 Synapse数据集

3.2 ACDC数据集

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

论文PDF和代码下载

后台回复：0517，即可下载上述论文和代码

后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集

后台回复：Transformer综述，即可下载最新的两篇Transformer综述PDF

重磅！Transformer交流群成立

扫码添加CVer助手，可申请加入CVer-Transformer 微信交流群，方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。

一定要备注：研究方向+地点+学校/公司+昵称（如Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲长按加小助手微信，进交流群

▲点击上方卡片，关注CVer公众号

整理不易，请给CVer点赞和在看