Transformer再下一城!Swin-Unet:首個純Transformer的醫學圖像分割網絡


OpenCV學堂 今天
 

以下文章來源於集智書童 ,作者ChaucerG

集智書童
集智書童

機器學習知識點總結、深度學習知識點總結以及相關垂直領域的跟進,比如CV,NLP等方面的知識。

圖片

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等,代碼即將開源!
作者單位:慕尼黑工業大學, 復旦大學, 華為(田奇等人)

圖片

摘要

 

在過去的幾年中,卷積神經網絡(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和卷積的組合。

圖片

框架結構


圖片圖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.1 Synapse數據集

圖片 圖片

3.2 ACDC數據集

圖片

就很多基於Transformer的文章和已經開源的代碼來看,模型的預訓練會嚴重影響基於Transformer模型的性能。其實作者在文章的最后也提到了他們也是直接使用Swin-Transformer在ImageNet上的預訓練權值來初始化網絡的Encoder和Decoder的,這可能是一個次優方案。不過作者也在努力嘗試實現完全端到端的實驗了(此處點贊)。圖片


圖片


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM