Unet系列(1) — Unet++論文筆記
最初原文應為[1]，本文另參考[2]，作為補充。兩篇的第一作者為同一人。
Unet/FCN在semantic segmentation的侷限
首先，如果對FCN或Unet的架構不是非常了解的朋友，推薦可以先看這篇：
也是我的出道作啦，哈哈。Unet和FCN系列的特色有二：
- Upsampling，或者說decoder段。將萃取的特徵放回原本輸入圖片的大小，他想探究的是這些細部特徵在圖片上的位置。
- Skip connection。在upsampling的過程中將對應downsampling同一層級的特徵圖串接(concatenate)，藉以保留在不同層級萃取到的特徵。
這兩個特色讓Unet/FCN對於像素層級(pixel-level)的預測變得有力(robust)。然而，Unet++的作者也針對這兩點，提出他認為反而造成限制的地方：
- 不知道Unet/FCN最適合的深度是多少。也就是理想的層數。這個沒有辦法用先驗的方式，而必須使用額外的架構，或者就訓練不同層數的Unet，然後再用集成(ensemble)的方式作推論(inference)。然而這樣做將相當耗費時間和資源，降低效率。
- Skip connection要求將對應層級的特徵圖串接在一起，然而這也是個限制：「只能」串接同一層級的特徵圖，在Unet架構裡，通常只有一張。
Unet++的解決方法
- 將好幾個不同深度的Unet放在一起。
- 這些Unet之間資訊流通，讓因為進到不同深度的Unet，其upsampling不同程度(scale)的特徵圖也都可以被使用到。
作者期待：
- 提高效率
- 模型更有彈性
- 當然，最重要的是表現要更好
那具體是怎麼做的呢?我們繼續往下討論。
架構
Unet的一層架構如下：
X⁰⁰到X¹⁰是encoder，下降段。X¹⁰到X⁰¹是decoder，上升段。X⁰⁰到X⁰¹是skip connection。這樣的狀況，拉到四層深的時候是這樣：
現在，為了同時訓練多個Unet，作者的想法很簡單：把不同深度的decoder疊進去。
這些Unet群共用一個encoder，而各自有decoder。為了能同時訓練，其loss被集合起來計算，然後再各自反向傳播梯度回去。關於損失函數的部分，下面會再試著說明更仔細一點。
接著，為了讓不同深度的Unet之間的資訊也可以流通，把connection做一點小小的調整：
我們以第一層為例說明。原本共用一個encoder，所以在skip connection的時候，都是由X⁰⁰傳出去，分別給X⁰¹、X⁰²、X⁰³、X⁰⁴。現在為了讓不同深度的Unet可以流通，X⁰¹的資訊也要給X⁰²，X⁰²的資訊要給X⁰³，以此類推。你一定會想：那這樣原本的skip connection就消失了耶?沒錯，所以作者最後要把他們混~合~在~一~起
要注意的是，skip connection不僅僅在encoder/decoder之間，連不同的decoder之間也都有。如此一來，對於某個decoder單元而言，只要它不是位在最淺的那個Unet，它都可以收到
- 來自下一層，同一個decoder的單元
- 同一層，其他比它更淺的Unet的decoder的單元
- 同一層encoder的單元
的訊息。寫成式子就是：
例如X⁰³：
同時，它也要傳出給同一層，比它更深的decoder(除非它最深)，同時往同一個decoder的更上層傳遞(除非它最上層)。
這個架構可以達到作者想要的同時訓練不同深度Unet的目標，以及讓模型更有彈性。怎麼說?如果訓練完整個Unet++之後，發現架構太厚重了，想拿掉一部分，我可以：
根據自己的喜好隨意調整模型。當然，這可能會減損部分模型表現。
深層監督
簡而言之，作者在損失函數做了一些改變。他希望兼顧梯度下降的平滑性和克服面對一般圈選任務最常遇到的問題：class imbalance，所以將cross entropy和dice coefficient結合在一起：
C是指分類項目數，N則是一個批次(batch)所有樣本(sample)的像素數目總合。y是標註結果(ground truth)，而p則是在該畫素，模型判斷各分類的結果，以機率表現。這個模型判斷的結果是針對模型中的一個單元(以上面的圖來說，就是X⁰¹、X⁰²、X⁰³、X⁰⁴之一)，而要達到深層監督，還有一個步驟要做，就是加權總合(weight sum)。不過文章中作者也很坦白說，他們並未設立權重，因此只是加總而已，每個單元的輸出是一樣比重的。
另外，在從單元輸出之前，他們會用一個1×1的卷積網路讓輸出符合類別數(所以會有C個卷積核)，以及經過sigmoid啟動函數，不過並未說明設定的閾值，所以不確定是不是只是經過sigmoid處理而已。
實驗結果
和Unet比較
兩篇文章都把重點擺在生醫相關的資料集。例如，在[2]，作者嘗試了細胞、電子顯微鏡影像、腦部、肺部，及肝臟的影像。直接把結果貼在這邊：
另外放上客觀數據結果：
從文章發表的內容看起來，作者也達到了「表現要更好」的目標。
修剪的效應
因為作者還強調模型的表現不會受到「修剪」的影響太多，所以他們另外比較：把模型訓練好以後再修剪vs.單獨訓練特定深度的模型，兩者的表現：
看起來一起訓練是比較好的。L⁴因為是最深的那個，理論上不會修剪，因此一起訓練和單獨訓練是差不多的。
最後是比較隨著模型修剪，不同深度的模型的表現：
L¹是陽春而且最淺的Unet，只有一層而已，因此表現差也正常。其他三種看起來差異不大。
可視化(feature map visulization)
作者將不同的神經網路產生的特徵圖視覺化，提供讀者比對：
作者認為Unet的特徵圖跳太快了，直接從圖片角度跳到像素角度。利用重新建立的skip connection，特徵得以較為漸進的方式被萃取。我則認為深層監督對於較淺的模型的特徵萃取是有幫助的。
最後面看一下作者放的效率比較吧。平均而言，Unet++的收斂比Unet快一些。
結論
作者基於Unet本身的限制，提出叢集式的訓練，將不同深度的Unet疊在一起同時訓練，並重新設計skip connection。在損失函數，作者使用深層監督，有效幫助特徵的萃取。Unet++看起來達到了作者的要求：訓練有效率、可修剪以符合推論需求，表現也比較好。
實際上如何呢?期待下一篇實作篇吧!
參考資料
[1] Zhou et al., Unet++: A nested u-net architecture for medical image segmentation. DOI: 10.1007/978–3–030–00889–5_1
[2] Zhou et al., UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation. DOI: 10.1109/TMI.2019.2959609
