U-Net網絡詳解

U-Net: Convolutional Networks for Biomedical Image Segmentation

原文地址：https://zhuanlan.zhihu.com/p/43927696

前言

U-Net是比較早的使用全卷積網絡進行語義分割的算法之一，論文中使用包含壓縮路徑和擴展路徑的對稱U形結構在當時非常具有創新性，且一定程度上影響了後面若干個分割網絡的設計，該網絡的名字也是取自其U形形狀。

U-Net的實驗是一個比較簡單的ISBI cell tracking數據集，由於本身的任務比較簡單，U-Net緊緊通過30張圖片並輔以數據擴充策略便達到非常低的錯誤率，拿了當屆比賽的冠軍。

論文源碼已開源，可惜是基於MATLAB的Caffe版本。雖然已有各種開源工具的實現版本的U-Net算法陸續開源，但是它們絕大多數都刻意迴避了U-Net論文中的細節，雖然這些細節現在看起來已無關緊要甚至已被淘汰，但是爲了充分理解這個算法，筆者還是建議去閱讀作者的源碼，地址如下：https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/

1. 算法詳解

1.1 U-Net的網絡結構

直入主題，U-Net的U形結構如圖1所示。網絡是一個經典的全卷積網絡（即網絡中沒有全連接操作）。網絡的輸入是一張  的邊緣經過鏡像操作的圖片（input image tile），關於“鏡像操作“會在1.2節進行詳細分析，網絡的左側（紅色虛線）是由卷積和Max Pooling構成的一系列降採樣操作，論文中將這一部分叫做壓縮路徑（contracting path）。壓縮路徑由4個block組成，每個block使用了3個有效卷積和1個Max Pooling降採樣，每次降採樣之後Feature Map的個數乘2，因此有了圖中所示的Feature Map尺寸變化。最終得到了尺寸爲  的Feature Map。

網絡的右側部分(綠色虛線)在論文中叫做擴展路徑（expansive path）。同樣由4個block組成，每個block開始之前通過反捲積將Feature Map的尺寸乘2，同時將其個數減半（最後一層略有不同），然後和左側對稱的壓縮路徑的Feature Map合併，由於左側壓縮路徑和右側擴展路徑的Feature Map的尺寸不一樣，U-Net是通過將壓縮路徑的Feature Map裁剪到和擴展路徑相同尺寸的Feature Map進行歸一化的（即圖1中左側虛線部分）。擴展路徑的卷積操作依舊使用的是有效卷積操作，最終得到的Feature Map的尺寸是  。由於該任務是一個二分類任務，所以網絡有兩個輸出Feature Map。

圖1：U-Net網絡結構圖

如圖1中所示，網絡的輸入圖片的尺寸是  ，而輸出Feature Map的尺寸是 ，這兩個圖像的大小是不同的，無法直接計算損失函數，那麼U-Net是怎麼操作的呢？

1.2 U-Net究竟輸入了什麼

首先，數據集我們的原始圖像的尺寸都是  的。爲了能更好的處理圖像的邊界像素，U-Net使用了鏡像操作（Overlay-tile Strategy）來解決該問題。鏡像操作即是給輸入圖像加入一個對稱的邊（圖2），那麼邊的寬度是多少呢？一個比較好的策略是通過感受野確定。因爲有效卷積是會降低Feature Map分辨率的，但是我們希望  的圖像的邊界點能夠保留到最後一層Feature Map。所以我們需要通過加邊的操作增加圖像的分辨率，增加的尺寸即是感受野的大小，也就是說每條邊界增加感受野的一半作爲鏡像邊。

圖2：U-Net鏡像操作

根據圖1中所示的壓縮路徑的網絡架構，我們可以計算其感受野：

這也就是爲什麼U-Net的輸入數據是  的。572的卷積的另外一個好處是每次降採樣操作的Feature Map的尺寸都是偶數，這個值也是和網絡結構密切相關的。

1.3 U-Net的損失函數

ISBI數據集的一個非常嚴峻的挑戰是緊密相鄰的物體之間的分割問題。如圖3所示，(a)是輸入數據，(b)是Ground Truth，(c)是基於Ground Truth生成的分割掩碼，(d)是U-Net使用的用於分離邊界的損失權值。

圖3：ISBI數據集樣本示例

那麼該怎樣設計損失函數來讓模型有分離邊界的能力呢？U-Net使用的是帶邊界權值的損失函數：

其中  是$$softmax$$損失函數，  是像素點的標籤值，  是像素點的權值，目的是爲了給圖像中貼近邊界點的像素更高的權值。

其中  是平衡類別比例的權值，  是像素點到距離其最近的細胞的距離，  則是像素點到距離其第二近的細胞的距離。  和  是常數值，在實驗中  ，  。

2. 數據擴充

由於訓練集只有30張訓練樣本，作者使用了數據擴充的方法增加了樣本數量。並且作者指出任意的彈性形變對訓練非常有幫助。

3. 總結

U-Net是比較早的使用多尺度特徵進行語義分割任務的算法之一，其U形結構也啓發了後面很多算法。但其也有幾個缺點：

1. 有效卷積增加了模型設計的難度和普適性；目前很多算法直接採用了same卷積，這樣也可以免去Feature Map合併之前的裁邊操作
2. 其通過裁邊的形式和Feature Map並不是對稱的，個人感覺採用雙線性插值的效果應該會更好。

Reference

[1] Ronneberger O, Fischer P, Brox T. U-net: Convolutional networks for biomedical image segmentation[C]//International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015: 234-241.