Transformer又一城 | Swin-Unet:首個(gè)純Transformer的醫(yī)學(xué)圖像分割模型解讀


首個(gè)基于純Transformer的U-Net形的醫(yī)學(xué)圖像分割網(wǎng)絡(luò),其中利用Swin Transformer構(gòu)建encoder、bottleneck和decoder,表現(xiàn)SOTA!性能優(yōu)于TransUnet、Att-UNet等,代碼即將開(kāi)源!
作者單位:慕尼黑工業(yè)大學(xué), 復(fù)旦大學(xué), 華為(田奇等人)
1簡(jiǎn)介
在過(guò)去的幾年中,卷積神經(jīng)網(wǎng)絡(luò)(CNN)在醫(yī)學(xué)圖像分析中取得了里程碑式的進(jìn)展。尤其是,基于U形結(jié)構(gòu)和skip-connections的深度神經(jīng)網(wǎng)絡(luò)已廣泛應(yīng)用于各種醫(yī)學(xué)圖像任務(wù)中。但是,盡管CNN取得了出色的性能,但是由于卷積操作的局限性,它無(wú)法很好地學(xué)習(xí)全局和遠(yuǎn)程語(yǔ)義信息交互。
在本文中,作者提出了Swin-Unet,它是用于醫(yī)學(xué)圖像分割的類似Unet的純Transformer模型。標(biāo)記化的圖像塊通過(guò)跳躍連接被送到基于Transformer的U形Encoder-Decoder架構(gòu)中,以進(jìn)行局部和全局語(yǔ)義特征學(xué)習(xí)。
具體來(lái)說(shuō),使用帶有偏移窗口的分層Swin Transformer作為編碼器來(lái)提取上下文特征。并設(shè)計(jì)了一個(gè)symmetric Swin Transformer-based decoder with patch expanding layer來(lái)執(zhí)行上采樣操作,以恢復(fù)特征圖的空間分辨率。在對(duì)輸入和輸出進(jìn)行4倍的下采樣和上采樣的情況下,對(duì)多器官和心臟分割任務(wù)進(jìn)行的實(shí)驗(yàn)表明,基于純Transformer的U-shaped Encoder-Decoder優(yōu)于那些全卷積或者Transformer和卷積的組合。
2Swin-Unet架構(gòu)

Encoder, Bottleneck以及Decoder都是基于Swin-Transformer block構(gòu)造的實(shí)現(xiàn)。
2.1 Swin Transformer block

與傳統(tǒng)的multi-head self attention(MSA)模塊不同,Swin Transformer是基于平移窗口構(gòu)造的。在圖2中,給出了2個(gè)連續(xù)的Swin Transformer Block。每個(gè)Swin Transformer由LayerNorm(LN)層、multi-head self attention、residual connection和2個(gè)具有GELU的MLP組成。在2個(gè)連續(xù)的Transformer模塊中分別采用了windowbased multi-head self attention(W-MSA)模塊和 shifted window-based multi-head self attention (SW-MSA)模塊?;谶@種窗口劃分機(jī)制的連續(xù)Swin Transformer Block可表示為:

其中,和分別表示(S)W-MSA模塊和第塊的MLP模塊的輸出。
與前面的研究ViT類似,self attention的計(jì)算方法如下:

其中,表示query、key和value矩陣。和分別表示窗口中patch的數(shù)量和query或key的維度。value來(lái)自偏置矩陣
。2.2 Encoder
在Encoder中,將分辨率為的維tokenized inputs輸入到連續(xù)的2個(gè)Swin Transformer塊中進(jìn)行表示學(xué)習(xí),特征維度和分辨率保持不變。同時(shí),patch merge layer會(huì)減少Token的數(shù)量(2×downsampling),將特征維數(shù)增加到2×原始維數(shù)。此過(guò)程將在Encoder中重復(fù)3次。

Patch merging layer
輸入patch分為4部分,通過(guò)Patch merging layer連接在一起。這樣的處理會(huì)使特征分辨率下降2倍。并且,由于拼接操作的結(jié)果是特征維數(shù)增加了4倍,因此在拼接的特征上加一個(gè)線性層,將特征維數(shù)統(tǒng)一為原始維數(shù)的2倍。
2.3 Decoder
與Encoder相對(duì)應(yīng)的是基于Swin Transformer block的Symmetric Decoder。為此,與編碼器中使用的patch merge層不同,我們?cè)诮獯a器中使用patch expand層對(duì)提取的深度特征進(jìn)行上采樣。patch expansion layer將相鄰維度的特征圖重塑為更高分辨率的特征圖(2×上采樣),并相應(yīng)地將特征維數(shù)減半。

Patch expanding layer
以第1個(gè)Patch expanding layer為例,在上采樣之前,對(duì)輸入特征加一個(gè)線性層,將特征維數(shù)增加到原始維數(shù)的2倍。然后,利用rearrange operation將輸入特征的分辨率擴(kuò)大到輸入分辨率的2倍,將特征維數(shù)降低到輸入維數(shù)的1/4,即。
Up-Sampling會(huì)帶來(lái)什么影響?
針對(duì)Encoder中的patch merge層,作者在Decoder中專門設(shè)計(jì)了Patch expanding layer,用于上采樣和特征維數(shù)增加。為了探索所提出Patch expanding layer的有效性,作者在Synapse數(shù)據(jù)集上進(jìn)行了雙線性插值、轉(zhuǎn)置卷積和Patch expanding layer的Swin-Unet實(shí)驗(yàn)。實(shí)驗(yàn)結(jié)果表明,本文提出的Swin-Unet結(jié)合Patch expanding layer可以獲得更好的分割精度。

2.4 Bottleneck
由于Transformer太深導(dǎo)致收斂比較困難,因此使用2個(gè)連續(xù)Swin Transformer blocks來(lái)構(gòu)造Bottleneck以學(xué)習(xí)深度特征表示。在Bottleneck處,特征維度和分辨率保持不變。
2.5 Skip connection
與U-Net類似,Skip connection用于融合來(lái)自Encoder的多尺度特征與上采樣特征。這里將淺層特征和深層特征連接在一起,以減少降采樣帶來(lái)的空間信息損失。然后是一個(gè)線性層,連接特征尺寸保持與上采樣特征的尺寸相同。
skip connections數(shù)量的影響?
Swin-UNet在1/4,1/8和1/16的降采樣尺度上添加了skip connections。通過(guò)將skip connections數(shù)分別更改為0、1、2和3,實(shí)驗(yàn)了不同skip connections數(shù)量對(duì)模型分割性能的影響。從下表中可以看出,模型的性能隨著skip connections數(shù)的增加而提高。因此,為了使模型更加魯棒,本工作中設(shè)置skip connections數(shù)為3。

3實(shí)驗(yàn)結(jié)果
3.1 Synapse數(shù)據(jù)集


3.2 ACDC數(shù)據(jù)集

就很多基于Transformer的文章和已經(jīng)開(kāi)源的代碼來(lái)看,模型的預(yù)訓(xùn)練會(huì)嚴(yán)重影響基于Transformer模型的性能。其實(shí)作者在文章的最后也提到了他們也是直接使用Swin-Transformer在ImageNet上的預(yù)訓(xùn)練權(quán)值來(lái)初始化網(wǎng)絡(luò)的Encoder和Decoder的,這可能是一個(gè)次優(yōu)方案。不過(guò)作者也在努力嘗試實(shí)現(xiàn)完全端到端的實(shí)驗(yàn)了(此處點(diǎn)贊)。
4推薦閱讀

超越SOLO V2| ISTR:基于Transformer的端到端實(shí)例分割(文末獲取論文與源碼)

CVPR2021 | 重新思考BiSeNet讓語(yǔ)義分割模型速度起飛(文末獲取論文)

YOLO在升級(jí) | PP-YOLO v2開(kāi)源致敬YOLOV4攜帶Tricks又準(zhǔn)又快地歸來(lái)(附論文與源碼)

簡(jiǎn)單有效 | Transformer通過(guò)剪枝降低FLOPs以走向部署(文末獲取論文)

最快ViT | FaceBook提出LeViT,0.077ms的單圖處理速度卻擁有ResNet50的精度(文末附論文與源碼)
5參考
[1].Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation
[2].https://github.com/HuCaoFighting/Swin-Unet
本文論文原文獲取方式,掃描下方二維碼
回復(fù)【Swin-Unet】即可獲取論文
長(zhǎng)按掃描下方二維碼添加小助手。
可以一起討論遇到的問(wèn)題
聲明:轉(zhuǎn)載請(qǐng)說(shuō)明出處
掃描下方二維碼關(guān)注【集智書(shū)童】公眾號(hào),獲取更多實(shí)踐項(xiàng)目源碼和論文解讀,非常期待你我的相遇,讓我們以夢(mèng)為馬,砥礪前行!

