Swin-UNet：基于純 Transformer 結(jié)構(gòu)的語(yǔ)義分割網(wǎng)絡(luò)

Visual Transformer

Author：louwill

Machine Learning Lab

    

自從Transformer被引入計(jì)算機(jī)視覺(jué)以來(lái)，催生了大量相關(guān)研究與應(yīng)用。在圖像分割方向，涌現(xiàn)了像SETR和TransUNet等基于Transformer的語(yǔ)義分割網(wǎng)絡(luò)模型。

在TransUNet中，雖然引入了Transformer用于UNet編碼器，但其特點(diǎn)還是CNN與Transformer的混合編碼，解碼上也是基于CNN的上采樣。直觀上看，這種混合編碼的結(jié)構(gòu)并沒(méi)有完全發(fā)揮出Transformer的優(yōu)勢(shì)，并且作為backbone的ViT結(jié)構(gòu)也需要進(jìn)一步改進(jìn)。

而此前由MSRA提出的Swin Transformer正好作為視覺(jué)Transformer領(lǐng)域新的backbone。相較于TransUNet，去掉CNN編碼，用 Swin Transformer來(lái)代替原先的ViT，將UNet全部結(jié)構(gòu)都換成Swin Transformer。因而，基于Swin Transformer的Swin-UNet就應(yīng)運(yùn)而生。

提出Swin-UNet的論文為Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation，由慕尼黑工業(yè)大學(xué)、復(fù)旦和華為于2021年5月提出。

Swin-UNet結(jié)構(gòu)

Swin-UNet模型整體結(jié)構(gòu)如圖1所示。

如圖1所示，Swin-UNet由Encoder、Bottleneck、Decoder和跳躍連接組成。先看編碼器部分，輸入圖像先進(jìn)行patch partition，每個(gè)patch大小為4x4，輸入維度為H/4 x W/4 x 48，經(jīng)過(guò)linear embedding和兩個(gè)Swin Transformer block后特征圖尺寸為H/4 x W/4 x C，然后通過(guò)patch merging進(jìn)行下采樣，再經(jīng)過(guò)兩個(gè)Swin Transformer block后特征圖尺寸變?yōu)镠/8 x W/8 x 2C，最后再進(jìn)行一次同樣的下采樣操作即可完成編碼器的操作。可以看到，Swin-UNet編碼器每次按照2倍來(lái)縮小patch的數(shù)量，然后按照3倍來(lái)擴(kuò)大特征維度的數(shù)量。

Bottleneck則是用了兩個(gè)連續(xù)的Swin Transformer block，這里為防止網(wǎng)絡(luò)太深不能收斂，所以只用了兩個(gè)block，在Bottleneck中，特征尺寸保持H/32 x W/32 x 8C不變。

然后是解碼器部分。Swin-UNet解碼器主要由patch expanding來(lái)實(shí)現(xiàn)上采樣，作為一個(gè)完全對(duì)稱的網(wǎng)絡(luò)結(jié)構(gòu)，解碼器也是每次擴(kuò)大2倍進(jìn)行上采樣，核心模塊由Swin Transformer block和patch expanding組成。

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

一個(gè)Swin Transformer block由一個(gè)W-MSA和一個(gè)SW-MSA組成，如圖2所示。

Swin-UNet實(shí)驗(yàn)

跟TransUNet一樣，Swin-UNet分別在Synapse多器官分割數(shù)據(jù)集和ACDC (自動(dòng)化心臟診斷挑戰(zhàn)賽)上實(shí)驗(yàn)了效果。在Synapse多器官分割上的效果如下表所示。

可以看到，相較于TransUNet等其他網(wǎng)絡(luò)，Swin-UNet綜合分割效果是最好的。可視化的分割效果如圖3所示。

除此之外，作者也做了一些消融研究 ，探討了不同的上采樣策略、跳躍連接的數(shù)量、輸入圖像尺寸、模型大小等多個(gè)方面探討了Swin-UNet的穩(wěn)健性。

Swin-UNet官方代碼已開(kāi)源：

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

總結(jié)

總體來(lái)看，Swin-UNet還是得益于Swin Transformer這樣一個(gè)強(qiáng)大的視覺(jué)Transformer backbone。作為第一個(gè)純Transformer結(jié)構(gòu)的語(yǔ)義分割模型，Swin-UNet值得大家一試。

往期精彩：
 Swin Transformer：基于Shifted Windows的層次化視覺(jué)Transformer設(shè)計(jì)
 TransUNet：基于 Transformer 和 CNN 的混合編碼網(wǎng)絡(luò)
 SETR：基于視覺(jué) Transformer 的語(yǔ)義分割模型
 ViT：視覺(jué)Transformer backbone網(wǎng)絡(luò)ViT論文與代碼詳解
【原創(chuàng)首發(fā)】機(jī)器學(xué)習(xí)公式推導(dǎo)與代碼實(shí)現(xiàn)30講.pdf
【原創(chuàng)首發(fā)】深度學(xué)習(xí)語(yǔ)義分割理論與實(shí)戰(zhàn)指南.pdf
求個(gè)在看