3D Unet 小簡介:架構(1) — 從2D開始

Martin Huang
11 min readJul 23, 2020

前言

在前輩的鼓勵以及要求之下,我決定來嘗試看看這種技術型筆記。先起個頭,相信後面會越來越習慣做類似的事情,也可以把一些經驗或知識記錄在這邊,方便往後查閱。同時又能整理已經學過的東西,或許可以在之中找到突破口。

在進入文章整理之前,容我再敘述一小段選擇這個神經網路的動機。從客觀條件來看,我的準備並不算充分,無論是對程式語言的掌握度,或者對深度學習的了解,都是如此。我不敢說我已經駕馭,或至少游刃有餘的用程式工具解決事情;對於深度學習,無論是整體脈絡或者僅就Unet而言,也不敢說有掌握到精髓。學習到今日,憑藉的說穿了就是跨領域的初生之犢,憨膽,以及自身的興趣,還有來自於工作場所的壓力,如此而已。

我的工作領域是醫學。時至今日,神經網路甚至深度學習的相關技術已經和這個領域的諸多次領域有密切的結合,無論是在電腦視覺或是自然語言都是。各家醫院無不宣導所謂人工智慧、深度學習等等,舉辦大量座談會、研討會、視訊,邀請廠商、研究人員、臨床工作者參與。相關的產品也早有成果研製,放眼國內外許多與醫學領域結合的深度學習論文也層出不窮。但其實這兩者之間要如何結合,以及是否確實結合,在我看來其實還有很大的發展空間,或許在之後的文章有機會可以討論,但至少不是現在。

其實就算是與醫學結合,也不一定要選擇Unet來做,有很多其他相對簡單的模型可以用來做其他很多的工作,也比較適合技術還不成熟的我。可惜當初我是先決定題目的。照理來說,我應該直接找相關的人才來協助或合作,然後繼續往我的題目走。從臨床的角度來看,我選擇了基於醫療專業,目前現場尚未解決的問題,而我有興趣解決他。相關所需的技術我該找人。可惜,難找(不是人難找)。所以我決定自己跳下去做,然後發現,要選用的模型其實不簡單。

不相信的人也可以當作我在尋找藉口,但廢話說完了。這篇文章我要先整理Unet的模型,並在後面幾篇陸續介紹實作,以及目前遇到的問題。歡迎任何批評與指教。那就開始吧。

從2D Unet開始說起

直接上結構圖。原圖來自Olaf et al. U-Net: Convolutional Networks for Biomedical Image Segmentation. 發表時間為2015年8月。

Figure 1. Unet結構

其命名原因可以從結構的形狀得知。不過在專注於Unet的架構之前,還是必須要大約地介紹一下脈絡,因為它是承先,也是啟後。

物件辨識算是電腦視覺的一個延伸,從早期的神經網路以辨識/分類圖片為任務目標,慢慢發展到希望能將目標在圖片定位。其核心概念就是「將標籤推廣到畫素的層級」,從一整張圖用一張標籤,到每一個畫素都賦予它標籤。物件辨識裡還分有辨識(object detection)和圈選(segmentation)。兩者也都可以使用在醫學上,例如病灶的位置,或者重要器官等等。

例如在Unet文章中,提到2012年Ciresan提出的文章[1],使用的是sliding window的概念來辨識細胞膜構造;但以當時的狀況而言,經由sliding window裁切的小圖片雖然可以增加訓練量(變相的data augmentation),但也增加訓練時間,而且有許多裁切的部分重疊,其實還浪費訓練資源。此外,大的裁切圖片需要經過大的池化層而造成定位準確度下降,可是小的裁切圖片則含括太少資訊。

後續則有使用"hypercolumn"擷取單一畫素在每一層的活動度(vector of activation)做相加[2]以及”caseded hierarchical model”[3]利用萃取降維的過程中學習,再利用升維將前面降維的學習結果和最後的結果相結合的方法。

Figure 2. Hypercolumn技巧
Figure 3. Cascaded Hierarchical Models看起來已經有Unet的雛型

行文至此我忍不住會想把Feature pyramid network拿來比較。雖然它是較晚的文章(2017),但是構型也有些類似。

Figure 4. Feature pyramid Network(FPN)

和Unet的差異大概是Unet主要是在最後一層才做prediction(top-down),而FPN則每層都做,事實上關注的重點在於如何從不同規格大小的圖片擷取特徵。

Figure 5. top-down和FPN

Unet的架構基本上是從fully convolutional network(FCN)演化而來,這個架構的重點文章最早是2014年提出的[4]。有幾個特點:
1. 全架構卷積化(convolutionalization)。將以往神經網路最後幾層的全連結層(full connection layer)換成卷積網路,目的是為了使不同大小的圖片都可以輸入,而不必被限制。或者說,若把全連結層視為1*1的卷積網路,那就是把它擴張。
2. Upsampling。因為池化層會使圖片降維,若要將萃取到的特徵以「熱圖」(heatmap)的方式和原圖對照以找出物件,則必須重新拉回到原圖大小。為此,作者試了三個方法,但礙於篇幅且FCN並非本篇主角,在此簡單提過:
(1) Shift-and-stich。
(2) Bilinear interpolation。
(3) Backwards convolution(有時被稱作deconvolution,但如此稱呼稍嫌不精確)。
作者最後使用(3),原因是快速且有效。
3. Skip architecture。神經網路末端的特徵圖雖然較細緻但無法綜覽全圖,而前端的特徵圖雖然範圍大但特徵較粗糙,因此最好的方法是兩者都要學習,將其結合起來。以此概念設計一個DAG(有向無環圖),得到的結論是將進到下一層的卷積時保留本層的特徵並做預測,再和之後做convolution時得到的特徵圖相拼。由於兩者維度不同,因此必須先經過deconvolution之後才有辦法相合。

Figure 6. FCN 結構及預測。 FCN-32是只有最後的feature map做預測;FCN-16則是加上pooling 4,FCN-8再加上pooling 3。疊加之後的確有越精細的趨勢。

Unet建構的目的是為了使用少量的資料就能進行預測。結構圖中左邊的降維路線稱為contracting path,右邊則為expanding path。在進入降維或升維之前,都會經過連續兩層的卷積,被稱為successive convolution,作者認為這樣可以讓神經網路更容易抓到精確的訊息。另外一個特色是維持upsampling的高channel數目,使相對位置關係和細部的特徵可以充分的結合,增加辨識的品質。另外就是神經網路可以自行以鏡像補充圖片,如此一來對於在圖片邊緣的物件辨識度就可以增加,且不占用GPU的消耗,這是很重要的。

說了這麼多終於要把結構本身的細節好好分析一下了。整個過程中的卷積核大小都是3*3。啟動函數使用ReLu,pooling和upsampling都是2*2。不同的是,每次pooling之後channel數量加倍,而每次upsampling之後則減半,但因為會和對側的特徵圖串接,所以串接後總channel數目還是一樣,然後隨著deconvolution再減少。

這裡文章中有提到一句"The cropping is necessary due to the loss of border pixels in every convolution"(p.4, 2. Network Architecture, 第11行)。我不太確定他的意思是說經過卷積後相素會損失,還是說還要做額外的剪裁。畢竟如果不做padding,本來經過卷積後的相素質就會減少。從結構圖上的相素質算起來,應該是沒有經過額外的裁剪。這段如果理解有誤還請指正。

訓練過程

採用隨機梯度下降法(SGD),以較高的動量(momentum)提升權重和誤差的更新。為了最大化運用GPU,將每個batch的圖片量減為1,並用大量的batch來訓練。然後出現了一個叫energy function的名詞,一開始有點霧煞煞,後來查到可能是指結構動力學,或熱力學中類似熵的概念,亦即系統會向趨於穩定發展。梯度會向低點移動,則損失會越少。換言之,可以視為損失函數的概念。從文中的下一句敘述"The energy function is computed by a pixel-wise soft-max over the final feature map combined with the cross entropy loss function"可以得到確認。此外從這一句中可以得知它是利用每個相素的softmax和cross-entrophy結合計算的。數學公式就不推導了。

實驗結果

作者提供三種方法評估。
1. Pixel error。預測錯誤的相素數/總相素數。這個方法很直觀,但也有可能過度敏感。
2. Warping error。用在當pixel error很大,但乍看之下分割的形狀及位置還算準確時使用。主要用來評估預測的形狀。
3. Rand error。簡單說用來評估兩個(預測/實際)所分割的相素,其疊合度。
以下則是結果和比較。

若用IOU來評估的話:以下是另外兩種細胞的資料集辨認競賽結果。

小結

基本上當初在挑選神經網路時,看到這個模型,就心想:就是它了!不做他想。它完美符合我的需求。剩下的就是找到原始碼,或者尋找網路上的代碼做參考,以及製作屬於自己目標使用的資料集。
本來還想在這邊跨度到3D Unet,但看來篇幅已經過長,我決定另外開一篇文章,討論3D Unet的架構差異,以及2D和3D之間的適用性比較。

謝謝大家看到這邊。

Reference

[1] Ciresan, D.C., Gambardella, L.M., Giusti, A., Schmidhuber, J.: Deep neural networks segment neuronal membranes in electron microscopy images. In: NIPS. pp. 2852{2860 (2012)
[2] Hariharan, B., Arbelez, P., Girshick, R., Malik, J.: Hypercolumns for object segmentation and fine-grained localization (2014)
[3] Seyedhosseini, M., Sajjadi, M., Tasdizen, T.: Image segmentation with cascaded hierarchical models and logistic disjunctive normal networks. In: Computer Vision (ICCV), 2013 IEEE International Conference on. pp. 2168{2175 (2013)
[4] J. Long, E. Shelhamer, T. Shelhamer: Fully Convolutional Networks for Semantic Segmentation (2015)
[5] https://blog.csdn.net/weixin_36474809/article/details/87931260

--

--

Martin Huang

崎嶇的發展 目前主攻CV,但正在往NLP的路上。 歡迎合作或聯絡:martin12345m@gmail.com