UNet3+閱讀筆記

原創

2020-06-02 06:18

Paper:https://arxiv.org/pdf/2004.08790
Code:https://github.com/ZJUGiveLab/UNet-Version

以前和大家分享過UNet和UNet++，UNet是一種編碼器解碼器結構的卷積神經網絡，常用於醫學圖像分割；而UNet++是在UNet的基礎上添加嵌套和密集跳過連接的卷積神經網絡。這次給大家分享UNet系列的一篇新作UNet3+，它使用全尺度的跳躍連接把不同尺度的特徵圖相融合，並且通過深度監督從多尺度聚合的特徵圖中學習特徵表達，還提出了一種混合損失函數，將分類任務和分割任務相結合可以增強器官邊界和減少非器官圖像的過度分割，從而獲得更準確的分割結果。
UNet3+網絡結構如下圖1所示，通過重新設計跳躍連接來融合多尺度特徵，並利用全面的深度監督，提供更少的參數，從而產生更準確的分割結果。
圖1 UNet3+網絡模型

全尺度連接
爲了彌補UNet和UNet++不能精確分割圖像中器官的位置和邊界，UNet3+中每一個解碼器都結合了全部編碼器的特徵，這些不同尺度的特徵能夠獲取細粒度的細節和粗粒度的語義。比如，下圖2展示了第三個解碼器 $X_{De}^3$ 的特徵圖構造，它不僅直接接收來自相同尺度編碼器層的特徵圖 $X_{En}^3$ ，還對上面兩層進行不同尺度的最大池化操作，將 $X_{En}^1$ 和 $X_{En}^2$ 進行下采樣，保證下采樣後特徵圖的尺寸和 $X_{En}^3$ 的大小一樣；同時還對下面兩層進行不同尺度的上採樣操作，將 $X_{En}^4$ 和 $X_{En}^5$ 進行上採樣，確保上採樣後的特徵圖的尺寸和 $X_{En}^3$ 的大小一樣。最後，作者發現採用64個3x3的卷積核進行卷積效果較好，那麼在編碼器和解碼器之間構建的全尺度連接採用的就是64個3x3的卷積核進行卷積，在對應的解碼器模塊進行連接時，先採用320個(64x5=320)3x3的卷積核進行卷積，然後在經過BN和ReLU層就得到 $X_{De}^3$ 。

圖2 $X_{En}^3$ 的全尺度連接特徵圖
全尺度監督
爲了從全尺度的聚合特徵圖中學習層次表達，UNet3+進一步採用了全尺度的深度監督。與UNet++對每個嵌套的子網絡進行監督不同的是，在UNet3+中每一個解碼器模塊都有一個輸出，與GT(ground truth)進行比較計算loss，從而實現全尺度的監督。具體的流程是，每一個解碼器模塊的最後一層被送入一個3x3的卷積層，然後是一個雙線性上採樣得到和GT一樣大小的特徵圖，再使用sigmoid函數得到分割結果。
爲了實現器官邊界的精確分割，作者提出了一個多尺度結構相似指數損失函數(MS-SSIM，如公式1所示)來給予模糊邊界更高的權重。因此，UNet3+將更加關注模糊邊界，區域分佈差異越大，損失函數值(MS-SSIM)越大。

其中，M表示尺度的總數量(爲了和UNet、UNet++保持一致，作者將M設置爲5)， ${μ}_{p}$ , ${μ}_{g}$ 和 ${σ}_{p}$ , ${σ}_{g}$ 分別是p、g的均值和方差， ${σ}_{pg}$ 表示它們的協方差。 ${β}_{m}$ , ${γ}_{m}$ 定義每個尺度中這兩個部分的相對重要性， ${C}_{1}=0.01^2$ 和 ${C}_{2}=0.03^2$ 是避免被0除。
結合Focal loss, MS-SSIM loss和IOU loss，作者提出一種混合損失函數用於三個不同層次的分割任務(像素分割、塊分割和圖像分割)中，它能獲取不同尺度目標的清晰的邊界。混合損失函數的公式如公式2所示。
分類引導模塊
爲了防止非器官圖像的過度分割，和提高模型的分割精度，作者通過添加一個額外的分類任務來預測輸入圖像是否有器官，從而實現更精準的分割。如圖3所示，對不同尺度的解碼器模塊的結果進行一系列的操作，包括Dropout，1x1的卷積，max pooling和 sigmoid函數，就可以得到一個分類結果(0表示沒有器官，1表示有器官)。然後，將分類結果和分割結果相乘，這樣就可以抑制對非器官圖像的過度分割，從而提高模型的分割精度，實現精準分割。

圖3 分類引導模塊示意圖
實驗結果
圖4展示了UNet、UNet++和UNet3+在肝臟和脾臟數據集上面的分割結果，我們可以看出UNet3+表現最好；圖5展示了UNet3+和5個SOTA分割模型在肝臟和脾臟數據集上面的分割結果，我們可以看出UNet3+也表現的很好。