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所示。

Swin-UNet实验

跟TransUNet一样，Swin-UNet分别在Synapse多器官分割数据集和ACDC (自动化心脏诊断挑战赛)上实验了效果。在Synapse多器官分割上的效果如下表所示。

可以看到，相较于TransUNet等其他网络，Swin-UNet综合分割效果是最好的。可视化的分割效果如图3所示。

除此之外，作者也做了一些消融研究 ，探讨了不同的上采样策略、跳跃连接的数量、输入图像尺寸、模型大小等多个方面探讨了Swin-UNet的稳健性。

Swin-UNet官方代码已开源：

https://github.com/HuCaoFighting/Swin-Unet

总结

总体来看，Swin-UNet还是得益于Swin Transformer这样一个强大的视觉Transformer backbone。作为第一个纯Transformer结构的语义分割模型，Swin-UNet值得大家一试。

往期精彩：

求个在看