1. 程式人生 > >U-Net:用於生物醫學影象分割的卷積網路

U-Net:用於生物醫學影象分割的卷積網路

轉自:

http://baijiahao.baidu.com/s?id=1600400677130320639&wfr=spider&for=pc

摘要

眾所周知,成功地訓練一個深度網路需要大量的已標註資料集。本文提出了一種新的網路結構和訓練方法,依靠資料增強將有效的標註資料更為有效地使用。該網路由收縮路徑(contracting path)和擴張路徑(expanding path)組成。其中,收縮路徑用於獲取上下文資訊(context),擴張路徑用於精確的定位(localization),且兩條路徑相互對稱。該網路能夠從極少影象端對端進行訓練,並且在ISBI競賽中,對於分割電子顯微鏡中的神經元結構的表現好於以前最好的方法(滑動視窗卷積網路)。除此之外,該網路的執行速度非常快,在最近的一個GPU上,分割一張512*512的影象不到一秒即可完成。

引言

在過去的兩年中,深度卷積網路在許多視覺識別任務中表現良好。雖然卷積網路已經存在了很長時間,但其受限於訓練集和網路結構的大小。直到Krizhevsky等人在ImageNet資料集上訓練了一個具有8層和數百萬個引數的大型網路,人們才開始訓練更大更深度的網路結構。

卷積網路常用於分類任務,其中影象的輸出為單一的類別標籤。然而,在許多視覺任務中,尤其是在生物醫學影象處理中,所期待的輸出結果中應包括如位置和每一畫素的類別標籤等資訊。而且,在生物醫學任務中通常難以獲取數以千計的訓練影象。因此,Ciresan等人在滑動視窗設定中訓練網路,以某一畫素的鄰域(patch)作為輸入,用於預測每個畫素的類別標籤。首先,該網路能夠實現定位操作;其次,對畫素的鄰域而言,訓練資料的數量需遠大於訓練影象的數量。在2012年的ISBI的EM分割競賽中,該網路成功摘奪桂冠。

上述網路中有兩點明顯的不足:

該網路執行效率很慢。對於每個鄰域,網路都要執行一次,且對於鄰域重疊部分,網路會進行重複運算;網路需要在精確的定位和獲取上下文資訊之間進行權衡。越大的patch需要越多的最大池化層,其會降低定位的精確度,而小的鄰域使得網路獲取較少的上下文資訊。

最近有人提出了一種分類器輸出的方法,其考慮了來自多層的特徵。這使得同時做到精確的定位和上下文的使用成為可能。

本文將建立一個更為優雅的網路結構,即全卷積網路。基於該網路進行相應的修改和擴充套件,使得網路能夠在較少的訓練影象的情況下執行,並做出更為精確的分割操作。如下圖所示,其展現了網路的整體結構。在 Fully Convolutional Networks for Semantic Segmentation 一文中,其主要思想為通過連續層補充通常的收縮網路,其中上取樣(upsample)操作取代了池化操作。因此,這些層提高了輸出結果的分別率。為了進行定位操作,將來自收縮路徑的高解析度特徵與上取樣相結合。隨後連續的卷積層可以基於該資訊學習,得出更為精確的輸出。



U-net網路結構圖

圖一:上圖為U-net網路結構圖(以最低分別率為32*32為例)。每個藍色框對應一個多通道特徵圖(map),其中通道數在框頂標,x-y的大小位於框的左下角;白色框表示複製的特徵圖;箭頭表示不同的操作。

在該網路中最重要的修改是在上取樣部分有大量的特徵通道,其允許網路將上下文資訊傳播至更高解析度的層。因而,擴張路徑或多或少與收縮路徑對稱,從而形成一個U型結構。該網路沒有全連線層,且僅使用每個卷積的有效部分,例如:分割圖僅包含在輸入影象中可獲得完整上下文資訊的畫素。該策略允許通過重疊-切片(overlap-tile)策略對任意大的影象進行無縫分割,如下圖所示。為了預測影象邊緣區域中的畫素的類別標籤,通過對輸入影象進行映象操作來推斷缺失的上下文資訊。這種切片策略對於將網路應用於大型影象是重要的,否則影象的解析度將受限於視訊記憶體(the GPU memory)。

重疊-切片策略示意圖

圖二:上圖為重疊-切片策略示意圖(此處為EM stacks中神經元結構的分割)。預測黃色方框中的分割需要藍色方框中的影象資料作為輸入,缺失的資料由映象推斷。

對於只要少量訓練資料可用的任務,對可用的訓練影象廣泛地應用彈性形變的影象增強方法。這使得網路能夠學習這種形變的不變性,無需提供帶註釋的影象語料庫(the annotated image corpus)供網路學習。這對於生物醫學影象的分割是特別重要的,因為形變是組織中最常見的變化,且實際的形變能有效地模擬。 Dosovitskiy等人在無監督學習中表明瞭學習不變性的重要性。

在許多細胞分割任務中的另外一個挑戰為對同一類別的接觸目標的分離,如下圖所示。為此,採用加權損失,即在細胞接觸之間的分離背景標籤中的損失函式中設定較大大的權重。

DIC顯微鏡記錄的HeLa細胞

圖三:上圖為DIC(微分干涉,differential interference contrast)顯微鏡記錄的HeLa細胞。圖a為原始影象;圖b為與真實的帶標記的分割資料相疊加,其中不同的顏色表示HeLa細胞的不同例項(instance);圖c為生成的分割掩膜,其中白色區域表示前景,黑色區域表示背景;圖d為用於網路學習邊緣畫素的帶加權損失的畫素圖。

U-net網路模型適用於各種生物醫學影象分割問題。本文將展示在EM stacks競賽中對神經元結構分割的結果,以及2015年ISBI的光學顯微鏡影象中的細胞分割結果。

網路結構

網路結構如圖一所示。U-net網路由一個收縮路徑(左邊)和一個擴張路徑(右邊)組成。其中,收縮路徑遵循典型的卷積網路結構,其由兩個重複的3*3卷積核(無填充卷積,unpadded convolution)組成,且均使用修正線性單元(rectified linear unit,ReLU)啟用函式和一個用於下采樣(downsample)的步長為2的2*2最大池化操作,以及在每一個下采樣的步驟中,特徵通道數量都加倍。在擴張路徑中,每一步都包含對特徵圖進行上取樣(upsample);然後用2*2的卷積核進行卷積運算(上卷積,up-convolution),用於減少一半的特徵通道數量;接著級聯收縮路徑中相應的裁剪後的特徵圖;再用兩個3*3的卷積核進行卷積運算,且均使用ReLU啟用函式。由於在每次卷積操作中,邊界畫素存在缺失問題,因此有必要對特徵圖進行裁剪。在最後一層,利用1*1的卷積核進行卷積運算,將每個64維的特徵向量對映網路的輸出層。總而言之,該網路有23個卷積層。

為了允許輸出圖的無縫拼接,如圖二所示,以及所有2*2最大池化操作能應用於x、y大小相等(even)的層,輸入切片(tile)大小的選擇是很重要的。

訓練

輸入影象和其相應的分割圖用於訓練由基於Caffe的隨機梯度下降演算法實現的網路。由於未填充的卷積其輸出影象的大小小於恆定邊界寬度的輸入。為了最小化開銷和最大限度地利用視訊記憶體,比起輸入一個較大的批次,更傾向於較大的輸入切片,因而將批次大小縮減為單張影象。相應地,通過使用高動量(high momentum,0.99)使得大量先前的訓練樣本在當前的優化步驟中更新。

能量函式(the energy function)是通過與交叉熵損失函式相結合的最終特徵圖,並利用畫素級的soft-max函式來計算的。soft-max函式定義為:

其中,a k (x)表示在第x畫素點上第k個特徵通道的啟用函式。K為類別數,p k (x)為近視最大函式(the approximated maximum-function)。

加權交叉熵定義為:

其中,l={1,...,K}表示每個畫素正確的標籤,權重函式w,且將權重函式定義為:

其中,w c 表示平衡類別頻率的權重圖,d 1 表示到最近細胞邊界的距離,d 2 表示到第二近細胞邊界的距離。基於經驗,將w 0 設為10,σ約為5畫素。

對於分割邊界操作,使用了形態學方法。

權重初始化基於標準差為√(2/N)的高斯分佈,其中N表示每個神經元的輸入結點數。

資料增強

在3*3的網格上使用隨機位移向量產生平滑形變,其中位移來自於10畫素標準差的高斯分佈,且通過雙三次插值法計算得出。在收縮路徑的末尾的drop-out層進一步暗示了資料增強。