U-Net: Convolutional Networks for Biomedical Image Segmentation
U-Net:用於生物醫學圖像分割的卷積網絡
摘要
要想成功地訓練一個深度網絡需要大量的數以千計的有標記的樣本,這已經成為了業內共識。在本文中,我們提出了一種網絡和相應的訓練策略,它依賴於強大的數據擴充技術,以更有效地使用可用的有標記的樣本。該體系結構由捕獲context的contracting路徑和實現精確定位的symmetric expanding路徑構成。我們表明這樣的網絡可以用非常少的圖像端到端地進行訓練,並且優於先前ISBI挑戰賽中對電子顯微鏡堆疊中的神經元結構進行分割的最佳方法(滑動窗口卷積網絡)。使用在透射光顯微鏡圖像(相位對比度和DIC)上訓練的相同網絡,我們在這些類別中大幅度地贏得了2015年ISBI細胞追蹤挑戰。而且,網絡速度很快。 在最近的GPU上,512x512圖像的分割不到一秒鍾就可以完成。完整的實施(基於Caffe)和訓練的網絡可在以下網站獲得http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net .
1.緒論
在過去兩年中,深度卷積網絡在許多視覺識別任務中表現優於現有技術,例如,[7,3]。雖然卷積網絡已經存在了很長時間[8],但由於可用訓練集的大小和所考慮網絡的規模,它們的成功受到限制。Krizhevsky等人的突破 [7]是由於對包含100萬個訓練圖像的ImageNet數據集用有8層和數百萬個參數的大型網絡進行有監督訓練。在那之后,更大更深的網絡也可以訓練了。
卷積網絡的典型用途是分類任務,其中圖像的輸出是單個類別標簽。然而,在許多視覺任務中,尤其是在生物醫學圖像處理中,期望的輸出應該包括定位,即,應該將類標簽分配給每個像素。此外,生物醫學任務中通常無法拿到數千個將標記分配給每個像素訓練圖像。因此,Ciresan等人。 [1]在滑動窗口設置中訓練網絡,通過在該像素周圍提供局部區域(patch)作為輸入來預測每個像素的類標簽。首先,這個網絡可以定位。 其次,以patches方式組織的訓練數據遠遠大於訓練圖像的數量。由此產生的網絡在ISBI 2012上大幅度贏得了EM細分挑戰。
顯然,Ciresan等人的策略[1]有兩個缺點。 首先,它非常慢,因為必須為每個patch丁單獨運行網絡,並且由於patches重疊而存在大量冗余。其次,在定位准確度和context的使用之間存在trade-off。較大的patches需要更多的最大池層這會降低定位精度,而小patches則只允許網絡看到很少的context。最近的方法[11,4]提出了一種分類器輸出,其考慮了來自多個層的特征。同時兼具良好的定位與context的使用成為了可能。
在本文中,我們建立了一個更優雅的網絡結構,即所謂的“完全卷積網絡”[9]。我們對這種架構進行了修改和擴展,使得它只需很少的訓練圖像就可以進行更精確的分割;見圖1。[9]中的主要思想是通過successive layers補充通常的contracting網絡,其中池化操作由上采樣操作替換。因此,這些層增加了輸出的分辨率。為了進行定位,來自contracting路徑的高分辨率特征與上采樣輸出相結合。然后,successive卷積層可以基於該信息學習組裝出更精確的輸出。我們的架構中的一個重要修改是在上采樣部分中我們還有大量的特征通道,這些通道允許網絡將context信息傳播到更高分辨率的層。因此,expanding路徑或多或少地與contracting路徑對稱,並產生U形結構。網絡沒有任何全連接層,並且僅使用每個卷積的有效部分,即,分割圖僅包含在輸入圖像中可獲得完整context的像素。該策略允許通過overlap-tile策略無縫分割任意大的圖像(參見圖2)。為了預測圖像邊界區域中的像素,通過鏡像輸入圖像來外推缺失的context。這種平鋪策略對於將網絡應用於大圖像很重要,否則分辨率將受到GPU內存的限制。
至於我們的任務,可用的訓練數據非常少,我們通過對可用的訓練圖像應用彈性變形來獲得大量擴增的數據。這允許網絡學習這種變形的不變性,而不需要在有標簽的圖像語料庫中看到這些變換。這在生物醫學分割中尤其重要,因為變形是組織中最常見的變化,並且可以有效地模擬真實的變形。Dosovitskiy等人已經在無監督特征學習的范圍內證明了學習數據擴增的不變性的價值 [2]。
許多細胞分割任務中的另一個挑戰是分離同一類的相連物體; 見圖3。為此,我們提出使用加權損失,其中相連細胞之間的分離背景標簽在損失函數中獲得大的權重。
由此產生的網絡適用於各種生物醫學分割問題。在本文中,我們展示了EM堆棧中神經元結構分段的結果(ISBI 2012開始的持續競爭),其中我們的表現優於Ciresan等人的網絡[1]。此外,我們在2015年ISBI細胞追蹤挑戰的光學顯微鏡圖像中顯示了細胞分割的結果。在這里,我們在兩個最具挑戰性的2D透射光數據集上獲得了巨大的優勢。
2. 網絡結構
網絡結構如圖1所示。它由contracting路徑(左側)和expanding路徑(右側)組成。contracting路徑遵循卷積網絡的典型架構。它包括重復應用的兩個3x3卷積(unpadded卷積),每個卷積后跟一個整流線性單元(ReLU)和一個步長為2的用於下采樣的2x2最大池化操作。在每個下采樣步驟中,我們將特征通道的數量加倍。expanding路徑中的每一步都包括對特征圖通過2x2卷積(“向上卷積”)進行上采樣,這將會把特征通道數量減半,然后與來自contracting路徑的相應的裁剪后的特征圖串聯,再經過兩個3x3 卷積,每個都跟着ReLU。由於每個卷積中邊界像素的丟失,裁剪是必要的。在最后一層,使用1x1卷積將每個64分量特征向量映射到所需數量的類。總的來說,網絡有23個卷積層。
為了實現輸出分割圖的無縫平鋪(參見圖2),選擇輸入圖塊大小非常重要,這樣所有2x2最大池操作都應用於具有偶數x和y尺寸的圖層。
3. 訓練
輸入圖像及其相應的分割圖用於訓練網絡,采用Caffe [6]的隨機梯度下降實現。由於unpadded的卷積,輸出圖像小於輸入圖像一定的邊界寬度。為了最大限度地減少開銷並最大限度地利用GPU內存,我們傾向於在大batch size的情況下使用大輸入tiles,從而將batch減少到單個圖像。因此,我們使用高動量(0.99),使得大量先前看到的訓練樣本確定當前優化步驟中的更新。能量函數通過將最終特征圖上的逐像素soft-max與交叉熵損失函數組合來計算。soft-max被定義為:
${p_k}(x) = \exp ({a_k}(x))/(\sum\nolimits_{{k^'} = 1}^K {\exp ({a_{{k^'}}}(x))} $
其中${a_k}(x)$為在像素位置$x \in \Omega $,$\Omega \subset {Z^2}$處特征通道$k$處的激活。$K$是類別數,${p_k}(x)$是近似最大函數,即對於有最大激活${a_k}(x)$的$k$,${p_k}(x) \approx 1$,對於其他的$k$,${p_k}(x) \approx 0$。然后交叉熵在每個位置懲罰${p_{l(x)}}(x)$與1的偏差,使用:$E = \sum\limits_{x \in \Omega } {w(x)\log ({p_{l(x)}}(x))} $
其中,$l:\Omega \to \{ 1,...,K\} $是每個像素的真實標簽,$w:\Omega \to IR$是權重圖,我們引入它來給一些像素在訓練期間分配更大的權重。
我們預先計算每個ground truth分割的權重圖,以補償訓練數據集中某個類別的不同像素頻率,並迫使網絡學習我們在相連細胞之間引入的小分離邊界(見圖3c 和d)。使用形態學運算來計算分離邊界。然后將權重圖計算為:
$w(x) = {w_c}(x) + {w_0} \cdot \exp ( - {{{{({d_1}(x) + {d_2}(x))}^2}} \over {2{\sigma ^2}}})$
其中,${w_c}:w \to IR$是用來平衡類別頻率的權重圖,${d_1}:w \to IR$代表到最近細胞邊界的距離,${d_2}:w \to IR$到第二個最近細胞邊界的距離。在我們的實驗中,我們設置w0 = 10和σ≈5像素。
在具有許多卷積層和通過網絡的不同路徑的深度網絡中,權重的良好初始化是極其重要的。否則,網絡的某些部分可能會過度激活,而其他部分則不會發揮作用。理想情況下,應調整初始權重,使得網絡中的每個特征圖具有近似單位方差。對於具有我們的架構(交替卷積和ReLU層)的網絡,這可以通過從標准偏差為$\sqrt {2/N} $的高斯分布繪制初始權重來實現,其中N表示一個神經元的傳入節點的數量[5]。例如, 對於前一層中的3x3卷積和64個特征通道,N = 9·64 = 576。
3.1. 數據擴增
當只有少數訓練樣本可用時,數據增加對於教授網絡所需的不變性和魯棒性屬性至關重要。在顯微圖像的情況下,我們主要需要移位和旋轉不變性以及對變形和灰度值變化的魯棒性。特別是訓練樣本的隨機彈性變形似乎是訓練具有極少有標簽圖像的分割網絡的關鍵概念。我們使用粗略3乘3網格上的隨機位移矢量生成平滑變形。位移是從具有10個像素標准偏差的高斯分布中采樣的。然后使用雙三次插值計算每像素位移。contracting路徑末端的droupout層執行進一步的隱式數據擴充。
4. 實驗
我們演示了u-net在三種不同分割任務中的應用。第一項任務是電子顯微鏡記錄中神經元結構的分割。圖2顯示了數據集和我們獲得的分割的示例。我們提供完整的結果作為補充材料。該數據集是由ISBI 2012啟動的EM分割挑戰[14]提供的,並且仍然對新的貢獻者開放。訓練數據是來自果蠅第一齡幼蟲腹側神經索(VNC)的連續切片透射電子顯微鏡的一組30個圖像(512×512像素)。每個圖像都帶有相應的完全注釋的ground truth分割圖,用於細胞(白色)和膜(黑色)。該測試集是公開的,但其f分割圖保密。可以通過將預測的膜概率圖發送給組織者來獲得評估。通過在10個不同的級別對圖進行閾值處理並計算“扭曲誤差”,“隨機誤差”和“像素誤差”[14]來完成評估。
u-net(輸入數據的7個旋轉版本的平均值)在沒有任何進一步預處理或后處理的情況下實現0.0003529的扭曲誤差(新的最佳分數,參見表1)和0.0382的隨機誤差。
這明顯優於Ciresan等人的滑動窗口卷積網絡結果[1],其最佳提交的扭曲誤差為0.000420,隨機誤差為0.0504。就rand誤差而言,該數據集上唯一性能更好的算法使用高度數據集特定的后處理方法,應用於Ciresan等人的概率圖[1]。
我們還將u-net應用於光學顯微圖像中的細胞分割任務。 這種分割任務是2014年和2015年ISBI小區跟蹤挑戰的一部分[10,13]。第一組數據“PhC-U373”在聚丙烯酰亞胺底物上含有膠質母細胞瘤 - 星形細胞瘤U373細胞,通過相差顯微鏡記錄(參見圖4a,b和Supp。材料)。它包含35個部分注釋的訓練圖像。 在這里,我們實現了92%的平均IOU(“交聯結合”),這明顯優於第二個最佳算法(83%)(見表2)。第二組數據“DIC-HeLa”是通過微分干涉對比(DIC)顯微鏡記錄的平板玻璃上的HeLa細胞(參見圖3,圖4c,d和Supp。材料)。它包含20個部分注釋的訓練圖像。 在這里,我們的平均IOU為77.5%,明顯優於第二好的算法,為46%。
5. 結論
u-net架構在非常不同的生物醫學分割應用程序上實現了非常好的性能。由於彈性變形的數據增強,它只需要非常少的注釋圖像,並且在NVidia Titan GPU(6 GB)上只有10小時的非常合理的訓練時間。我們提供完整的基於Caffe [6]的實施和訓練的網絡。我們確信u-net架構可以輕松應用於更多任務。