Swin-UNet：基于纯 Transformer 结构的语义分割网络-轻识

Visual Transformer

Author：louwill

Machine Learning Lab

自从Transformer被引入计算机视觉以来，催生了大量相关研究与应用。在图像分割方向，涌现了像SETR和TransUNet等基于Transformer的语义分割网络模型。

在TransUNet中，虽然引入了Transformer用于UNet编码器，但其特点还是CNN与Transformer的混合编码，解码上也是基于CNN的上采样。直观上看，这种混合编码的结构并没有完全发挥出Transformer的优势，并且作为backbone的ViT结构也需要进一步改进。

而此前由MSRA提出的Swin Transformer正好作为视觉Transformer领域新的backbone。相较于TransUNet，去掉CNN编码，用 Swin Transformer来代替原先的ViT，将UNet全部结构都换成Swin Transformer。因而，基于Swin Transformer的Swin-UNet就应运而生。

提出Swin-UNet的论文为Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation，由慕尼黑工业大学、复旦和华为于2021年5月提出。

Swin-UNet结构

Swin-UNet模型整体结构如图1所示。

如图1所示，Swin-UNet由Encoder、Bottleneck、Decoder和跳跃连接组成。先看编码器部分，输入图像先进行patch partition，每个patch大小为4x4，输入维度为H/4 x W/4 x 48，经过linear embedding和两个Swin Transformer block后特征图尺寸为H/4 x W/4 x C，然后通过patch merging进行下采样，再经过两个Swin Transformer block后特征图尺寸变为H/8 x W/8 x 2C，最后再进行一次同样的下采样操作即可完成编码器的操作。可以看到，Swin-UNet编码器每次按照2倍来缩小patch的数量，然后按照3倍来扩大特征维度的数量。

Bottleneck则是用了两个连续的Swin Transformer block，这里为防止网络太深不能收敛，所以只用了两个block，在Bottleneck中，特征尺寸保持H/32 x W/32 x 8C不变。

然后是解码器部分。Swin-UNet解码器主要由patch expanding来实现上采样，作为一个完全对称的网络结构，解码器也是每次扩大2倍进行上采样，核心模块由Swin Transformer block和patch expanding组成。

最后是跳跃连接。跳跃连接可以算是UNet的特色，Swin-UNet也自然不例外。

一个Swin Transformer block由一个W-MSA和一个SW-MSA组成，如图2所示。