UNet的深入研究

UNet論文: 地址

UNet源代碼: 地址

 

一、介紹

1.1 U-Net

很多分割網絡都是基於FCNs做改進,包括這一章要講的U-net。U-Net是一篇基本結構非常好的網絡,由於網絡結構像U型,所以叫Unet網絡

U-Net不止應用於圖像分割(尤其是醫學圖像分割),也用於自然圖像生成的任務,所以它幾乎是骨幹(backbone)網絡的標配

 

1.2 FCN和U-Net的區別

這裏我引用了U-Net++作者的總結

在計算機視覺領域,全卷積網絡(FCN)是比較有名的圖像分割網絡,醫學圖像處理方向,U-Net可以說是一個更加炙手可熱的網絡,基本上所有的分割問題,我們都會拿U-Net先看一下基本的結果,然後進行“魔改”。

U-Net和FCN非常的相似,U-Net比FCN稍晚提出來,但都發表在2015年,和FCN相比,U-Net的第一個特點是完全對稱,也就是左邊和右邊是很類似的,而FCN的decoder相對簡單,只用了一個deconvolution的操作,之後並沒有跟上卷積結構第二個區別就是skip connection,FCN用的是加操作(summation),U-Net用的是疊操作(concatenation)。這些都是細節,重點是它們的結構用了一個比較經典的思路,也就是編碼和解碼(encoder-decoder)結構,早在2006年就被Hinton大神提出來發表在了nature上.

當時這個encoder-decoder結構提出的主要作用並不是分割,而是壓縮圖像和去噪聲。輸入是一幅圖,經過下采樣的編碼,得到一串比原先圖像更小的特徵,相當於壓縮,然後再經過一個解碼,理想狀況就是能還原到原來的圖像。這樣的話我們存一幅圖的時候就只需要存一個特徵和一個解碼器即可。同理,這個思路也可以用在原圖像去噪,做法就是在訓練的階段在原圖人爲地加上噪聲,然後放到這個編碼解碼器中,目標是可以還原得到原圖。

後來把這個思路被用在了圖像分割的問題上,也就是現在我們看到的FCN或者U-Net結構,在它被提出的三年中,有很多很多的論文去講如何改進U-Net或者FCN,不過這個分割網絡的本質的拓撲結構是沒有改動的。舉例來說,去年ICCV上凱明大神提出的Mask RCNN. 相當於一個檢測,分類,分割的集大成者,我們仔細去看它的分割部分,其實使用的也就是這個簡單的FCN結構。說明了這種“U形”的編碼解碼結構確實非常的簡潔,並且最關鍵的一點是好用。

此外, 由於UNet也和FCN一樣, 是全卷積形式, 沒有全連接層(即沒有固定圖的尺寸),所以容易適應很多輸入尺寸大小,但並不是所有的尺寸都可以,需要根據網絡結構決定, 在U-Net中, 池化是2x2的且是valid策略,即沒有padding,因此要保證輸入的圖像在經過每一次池化的時候都要是邊長偶數。所以要特別注意輸入圖像的尺寸。一個比較好的方法是從最小分辨率(分辨率就是feature map的尺寸
)出發沿收縮路徑的反方向進行計算,得到輸入圖像的尺寸。

 

1.3 U-Net爲什麼在醫學圖像分割表現好?

U-Nnet模型在醫學影像分割中表現好, 並且在以後許多優秀的醫學圖像分割網絡中都是以它爲basebone. 爲了更好的分析原因, 不妨先看看醫學圖像有什麼特點吧

1.3.1 醫學影像的特點

尤其是相對於自然圖像而言,醫學圖像具備以下幾個特點:

  1. 圖像語義較爲簡單、結構較爲固定。我們做腦的,就用腦CT和腦MRI,做胸片的只用胸片CT,做眼底的只用眼底OCT,都是一個固定的器官的成像。由於器官本身結構固定和語義信息沒有特別豐富,所以高級語義信息和低級特徵都顯得很重要(UNet的skip connection和U型結構就派上了用場)。
  2. 數據量少。醫學影像的數據獲取相對難一些,很多比賽只提供不到100例數據。所以我們設計的模型不宜過大,如果參數過多,很容易導致過擬合。
  3. 多模態。相比自然影像,醫療影像比較有趣和不同的一點是,醫療影像是具有多種模態的。以ISLES腦梗競賽爲例,其官方提供了CBF,MTT,CBV,TMAX,CTP等多種模態的數據。

    這就需要我們更好的設計網絡去提取不同模態的特徵feature。這裏提供兩篇論文供參考: Joint Sequence Learning and Cross-Modality Convolution for 3D Biomedical Segmentation(CVPR 2017) Dense Multi-path U-Net for Ischemic Stroke Lesion Segmentation in Multiple Image Modalities.

  4. 可解釋性重要。由於醫療影像最終是輔助醫生的臨牀診斷,所以網絡告訴醫生一個3D的CT有沒有病是遠遠不夠的,醫生還要進一步的想知道,病竈在哪一層,在哪一層的哪個位置,分割出來了嗎,能求體積嘛?同時對於網絡給出的分類和分割等結果,醫生還想知道爲什麼,也就是說從醫生角度出發,更希望知道醫療AI的可解釋性如何

    解釋醫學影像以Attention及CAM(class-activation-map)等可視化方法爲主。

    https://www.zhihu.com/question/359161489?utm_source=qq

    https://www.zhihu.com/question/48224234

1.3.2 分析

  • 醫學圖像相比於普通圖像,其複雜度是非常高的,灰度範圍大,邊界不清晰等特點, 並且人體內部結構相對固定,分割目標在人體圖像中的分佈很具有規律,語義簡單明確且沒有特別豐富, 而UNet結合了下采樣時的低分辨率信息(提供物體類別識別依據)和上採樣時的高分辨率信息(提供精準分割定位依據),此外還通過skip connection填補底層信息以提高分割精度.
  • 由於醫學影像的數據獲取相對難一些,數據量少, 所以如果模型的參數過多,很容易導致過擬合, 而U-Net模型大小較小,參數較少,因此比較適合.

PS: 分辨率就是feature map的尺寸

 

二、U-Net網絡模型

Unet包括兩部分:

  1. 特徵提取部分,每經過一個池化層就一個尺度,包括原圖尺度一共有5個尺度。
  2. 上採樣部分,每上採樣一次,就和特徵提取部分對應的通道數相同尺度融合,但是融合之前要將其crop。這裏的融合也是拼接。

該網絡由收縮路徑(contracting path)和擴張路徑(expanding path)組成。其中,收縮路徑用於獲取上下文信息(context),即局部特徵,擴張路徑用於精確的定位(localization),且兩條路徑相互對稱。

 

2.1 編碼-解碼結構(encoder-decoder)

U-Net和FCN它們的結構都用了一個比較經典的思路,也就是編碼和解碼(encoder-decoder)結構, 這種結構本來不是用在圖像分割的,而是用在壓縮圖像和去噪聲的[在本篇中的1.2有提到]

2.1.1 編碼

前半部分是編碼, 它的作用是特徵提取(獲取局部特徵,並做圖片級分類),這個降採樣的理論意義是它可以增加對輸入圖像的一些小擾動的魯棒性,比如圖像平移,旋轉等,減少過擬合的風險,降低運算量,和增加感受野的大小

從神經網絡可視化那篇論文中《Visualizing and Understanding Convolutional Networks》能看出,底層的特徵(如Layer1,2,3的輸出)更偏向於組成圖像的基本單元,如點,線,邊緣輪廓,如下圖:

而在高層抽象的特徵(如Layer4,5)就更抽象,更近似於表示的是圖像的語義信息,更像是一個區域了。如下圖所示:

 

2.1.2 解碼

後半部分是解碼, 即利用前面編碼的抽象特徵來恢復到原圖尺寸的過程, 最終得到分割結果

那麼就有問題了,解碼在將下采樣(反捲積)數據恢復的時候,特徵scale會發生變化,必然會有信息的丟失,這個時候, skip connention的作用就凸顯出來了, skip connention起到了補充信息的作用.

 

2.2 skip connention 

skip connention 是爲了實現特徵融合的結構, 在FCN中用的是加操作(summation) ,而在U-Net中用的是疊操作(concatenation) ,這裏也指的是複製和剪切(copy and crop).總之skip connention起到了補充信息的作用,讓模型依賴更多的信息, 改善上採樣時信息不足的問題, 從而提高分割精度.

總結下, 語義分割網絡在特徵融合時有2種辦法:

  1. FCN式的逐點相加(即跳躍結構),對應caffe的EltwiseLayer層,對應tensorflow的tf.add()
  2. U-Net式的channel維度拼接融合(即copy and crop),對應caffe的ConcatLayer層,對應tensorflow的tf.concat()

在U-Net網絡中可以發現灰色箭頭的複製和剪切(copy and crop)操作,在同一層左邊的最後一層要比右邊的第一層要大一些,所以要想利用淺層的feature,就要進行一些剪切(以中心某個區域進行裁剪)。

 

此外, 可以看到, UNet融合了不同尺度的特徵,同時跳級連接保證上採樣恢復出來的特徵不會很粗糙, 一個經驗的解釋就是跳級連接能夠保證特徵更加精細。這裏推薦一篇論文,發表在AAAI 19 Workshop on Network Interpretability for Deep Learning的《Visualized Insights into the Optimization Landscape of Fully Convolutional Networks》,它通過可視化loss landscape對比發現這種跳級連接能夠促使網絡的minimizer更平坦,從而對新的數據敏感度更低,泛化能力更強

 

2.3 卷積

藍色箭頭代表3x3大小的卷積操作,並且stride是1,全程使用valid來進行卷積(卷積有三種模式,鏈接這裏),因此,每個該操作以後,feature map的大小會減2。

valid卷積,即僅使用每個卷積的有效部分,不填充

此外,關於卷積還有兩個重要的知識點:(具體可以參考我CNN的卷積知識)

  • 卷積層中的卷積核的通道只能設定爲該層輸入數據的通道,而大小可自定義(這裏是3x3)
  • 卷積層中的卷積核的數量==輸出特徵圖(output feature map)的數量(即輸出數據的通道數)

 

2.4 池化

紅色箭頭代表2x2的max pooling操作,需要注意的是,此時的padding策略也是vaild

 

2.5 反捲積操作

綠色箭頭代表2x2的反捲積操作,操作會將feature map的大小乘2。反捲積是上採樣的一種方式.

 

2.6 最後一層1X1卷積

最後一層使用1 X 1大小的卷積核,將通道數降低至特定的數量(如像素點的類別數量). 這個過程是如何運算的呢?首先我們知道卷積層中的卷積核的通道只能設定爲該層輸入數據的通道,而大小可自定義, 並且卷積層中的卷積核的數量==輸出特徵圖(output feature map)的數量(即輸出數據的通道數) ,所以這裏需要兩個大小爲1x1, 通道爲64的卷積核進行卷積運算,才能得到兩個通道的388*388大小的圖.如下面的兩個圖所示(左邊:多個卷積核的卷積運算模型圖, 來自這裏, 右變: 實際的卷積過程,步長S=1,填充=0)

最終的兩個通道分別表示背景和目標(前景)兩個類別

 

2.7 特徵圖上的像素級 softmax 值計算

接着模型的最後通過逐個像素地求其在這兩張圖像中該像素位置的最大數值描述(概率)作爲該像素的分類, 因此產生了一張已經分割好的圖片(output segmentation map),實現端到端.

那麼這個概率怎麼求, 論文中的損失函數首先是用了個pixel-wise softmax,就是每個像素對應的輸出單獨做softmax,也就是做了w*h個softmax, 

  • αk​(x) 代表在位置 x 處的像素在特徵圖中的第 k 層的激活值,也就是表示每一像素點(x)對應特徵通道(k)的得分
  • K 是像素點的類別總數
  • pk(x)是類k的對像素點x的分類結果,就是概率

 

2.8 加權LOSS

  • 是每個像素的真實標籤 
  • 是權重圖,用來讓某些像素更加重要 , 3.3會講怎麼獲得權重圖

 

三、創新(innovation)

下面的這些創新點並不適合所有任務,比如說對剛體進行分割,很難通過彈性變形對數據進行增強。

3.1 overlap-tile策略

在U-Net結構中,卷積全程都使用valid來進行卷積[2.3中提到],該卷積的特點就是沒有padding, 即padding = 0,所以特徵圖(feature map)會越卷越小,它不會越卷越大,導致了最後輸出的尺寸回不到原來輸入的尺寸

原本以爲U-Net是一個絕對對稱的結構, 然而並不是,它左右兩邊feature map的尺寸不一致,原因就是因爲valid卷積,只會越卷越小.如下圖的層A,如果這個卷積有padding,那麼它可以回到32x32,然後再上採樣2x2後,層B就變爲64x64,依次輪推最終實現結構的絕對對稱

顯然,如果生生把原始圖像喂進來,最後輸出的結果,會比原始圖像小好多。我們希望輸入跟輸出尺寸一樣,即不能強行scale up最後的輸出,又不想給每層卷積加padding(連續對feature map加padding卷積,會使得padding進來的feature誤差越來越大,因爲越卷積,feature的抽象程度越高,就更容易受到padding的影響,所以爲什麼要用valid卷積,這是我參考這個作者總結的看法)

於是就誕生了重疊-切片(overlap-tile)策略, 該策略的思想是:對圖像的某一塊像素點(黃框內部分)進行預測時,需要該圖像塊周圍的像素點(藍色框內)提供上下文信息(context),以獲得更準確的預測。簡單地說, 就是在預處理中,對輸入圖像進行padding, 通過padding擴大輸入圖像的尺寸,使得最後輸出的結果正好是原始圖像的尺寸, 同時, 輸入圖像塊(黃框)的邊界也獲得了上下文信息從而提高預測的精度

在U-Net原文中,用的鏡像操作(mirror padding),實際上可以使用別的padding方法, 比如在這裏做牙齒圖像的segmentation,使用到了zero padding,即擴大的尺寸全用0(黑色)來填充。如下圖所示:

因此, 使用Overlap-tile策略對數據進行預處理是有必要的,可對任意大的圖像進行無縫分割. 而這種方法通常需要將圖像進行分塊(patch)的時候才使用

那麼爲什麼要對圖像分塊(patch)再輸入,而不是直接輸入整張圖像呢?因爲內存限制,有的機器內存比較小,需要分塊。然後再通過Overlap-tile 策略將分塊圖像進行預處理擴大輸入尺寸.

總之,U-Net作者創新的Overlap-tile 策略,這種方法用於補全輸入圖像的上下信息,也可以解決由於現存不足造成的圖像輸入的問題

 

3.2 數據增強(data augmentation)

如果數據不夠的話,尤其是醫學影像的數據,可以用數據增強提高數據量.

U-Net作者的創新是使用了隨機彈性變形進行數據增強,對數據進行增強,增加數據量,這允許網絡可以學習到這種形變的不變形

此外,在收縮路徑的末尾的drop-out層進一步暗示了數據增強

 

3.3 加權loss

作者通過預先計算權重圖的來獲得每一個像素在損失函數中的權值,這種方法補償了訓練數據每類像素的不同頻率,並且使網絡更注重學習相互接觸的細胞間的邊緣。分割邊界使用形態學運算,特徵圖的計算方法如下:

  • ωc:Ω→R是用來平衡類頻率的權重圖 
  • d1:Ω→R表示到背景的某個像素點到最近細胞邊界的距離 
  • d2:Ω→R表示到背景的某個像素點到第二近的細胞邊界的距離

上圖是用DIC(二次干涉對比)顯微技術記錄的玻璃上的 HeLa 細胞。其中圖 (a) 是原始圖像(gt);圖 (b) 是基於 gt 的分割覆蓋。其中不同的顏色表示不同的 HeLa 細胞示例。圖 (c) 是生成的分割掩膜,其中白色部分是前景,黑色部分是背景;圖 (d) 是像素級損失權重圖,使得網絡強制學習邊緣像素。對於位與細胞接觸部分的像素加大權重(紅色的部分)

 

四、不足

U-Net++作者分析U-Net不足並如何做改進 :https://zhuanlan.zhihu.com/p/44958351

 

五、拓展

5.1 分塊處理

醫學圖像是一般相當大,但是分割時候不可能將原圖太小輸入網絡,所以必須切成一張一張的小patch,在切成小patch的時候,Unet由於網絡結構的原因,因此適合用overlap的切圖 [本篇的3.1 有講],overlap部分可以爲分割區域邊緣部分提供文理等信息, 並且分割結果並沒有受到切成小patch而造成分割情況不好。

 


參考文章:

https://zhuanlan.zhihu.com/p/31428783

https://blog.csdn.net/Formlsl/article/details/80373200

https://www.jianshu.com/p/14641b79a672

https://blog.csdn.net/hduxiejun/article/details/71107285

https://blog.csdn.net/qian99/article/details/85084686

https://www.zhihu.com/question/269914775?sort=created

https://blog.csdn.net/github_37973614/article/details/84559861

https://zhuanlan.zhihu.com/p/44958351

https://www.jianshu.com/p/b6a898753d29

https://www.zhihu.com/question/48224234

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章