3D Unet小簡介：架構(2) — 2D和3D

承接上一篇，當初雖然決定要做Unet，但一開始並未在2D與3D之間做選擇。回顧當時，直接決定做3D Unet似乎是太衝動了一點，沒有好好比較其中的差異。那為什麼當時會選擇3D?其實只是因為我的資料是電腦斷層影像(computed tomography，CT)，而它雖然是一張張的2D影像，但建構起來的是立體的概念。考量到標記會橫跨許多張切面，而且應該要建立z軸的順序及方位(context)，就選擇3D Unet。

事實上若以"medical image" “segmentation” “ Unet” 去尋找文獻，使用的模型五花八門，即使用"CT"搜尋，也只是提到Unet而已，沒有刻意提到2D或3D。其他還有與Unet架構相似的Vnet、Unet的改良型Unet++，甚至是ResNet結合Unet的ResUnet等等，各種神經網路不一而足[1]。本篇的重點還是想探討2D和3D的差異，以及若以立體醫學影像而言，其適用性如何。

以下的討論，將圍繞在Kayalıbay et al.[2]以及Isensee et al.[3]的文章討論3D Unet的架構，並參考網路其他資料[4–6]來比較適性。

3D Unet的架構

Kayalıbay提出的概念

從2D Unet擴展成3D主要是為了應付z軸的資訊。2D的神經網路也可以處理第三軸的資訊，但通常只能處理少量。或者說，如果資料的z軸資訊不多，可以把它當作2D來訓練。然而，面對立體影像，或如較大的器官(腦、肝臟)，其影像的切面將橫跨許多張，可能不是傳統2D神經網路可以負擔的。

在調整Unet使其能訓練3D影像時，最早的改變是Milletari引進residual block的概念[7]。residual block在ResNet被人熟知，原本設計的目的是為了減少因為堆疊較多的卷積層之後所造成的梯度消失以及飽和度問題。

Figure 1. Residual block

右圖是改良過的residual block。非線性的啟動函數在最後被移除使梯度的回饋更流暢。

無論是2D或3D Unet，訓練醫學影像最大的問題就是資料不足，因此必須活用資料增量的方法，旋轉、平移，改變圖片大小。另外還有一個"random deformation field"，讓每輪開始訓練圖片時使用不一樣的變形方式，使神經網路可以每輪都像是在訓練不同的圖片。

另一個問題是class imbalance。在相素為單位的標籤，大部分的相素都不是目標(都是背景)，所以目標的標籤相對於背景標籤的數量相差很大。如果在一張圖片中有複數個標籤，那這個情況會更嚴重。有不同的文章試著解決這個問題，例如Ronneberger[8]和Cicek[9]是使用weight map在損失函數上，而Milletari則選擇加大dice coefficient，這樣做的好處是不用增加超參數。

Figure 2. 3D Unet

基本上保留2D Unet的特色：有contracting path和expanding path，有skip connection。不同的點有二：一是在2D Unet的時候是把feature maps串接(concatenation)，而這裡則試著和相加(element-wise summation)比較。如此一來，就可以將contracting path的feature map及同一層expanding path上的feature map以下的區域視為一個大的residual block(ResNet是用summation)。

另外就是圖中右下方，將不同層的feature map也相加，這是參考了FCN的概念，可以讓分割更細緻(請看上一篇的圖)。相加的方法是：先把最低層的feature map做upsampling，再用雙線性內插(bilinear interpoloration)得到和第二低層一樣大小的圖，並與之相加。之後再upsampling，再取雙線性內插，然後跟第一層相加。但作者表示這樣的作法並不是為了讓分割更細緻，因為skip connection已經做到了；而是為了刺激較淺層的神經網路能輸出更好品質的feature map，從而讓訓練加速。

至於判讀3D影像時消耗大量記憶體的問題，一個常見的做法是將影像組拆成許多小片段，個別作判讀/分割之後再組合起來。另外一個方法則是在遇到記憶體負荷過大時，直接使用downsampling的結果，再用upsampling和內插的方式使其大小增加至和原圖相同。本文採用這個做法。

此外全架構使用padding以維持feature map的大小。feature maps的大小只在contracting或expanding的時候才改變。使用strided convolution取代max pooling，因為在之前的實驗中顯示這樣效果更好。啟動函數使用PReLu，conovlution的kernal大小為3*3*3，但最後輸出feature maps疊加以及deconvolution前的卷積則為1*1*1。

損失函數的dice coefficient稍微提一點好了。基本原理是評估預測和標籤的重合度，公式為

可以預想的是，其值落在0~1之間：如果兩者完全重合則為1，兩者完全無交集則為0。不過在文章中，他所用的其實是Jaccard Distance：

中間我跳過了一些東西。Union是指既在預測區，也在標籤區的相素。

作者將此結構拿去做兩個實驗：一個是手部MRI的辨識，另一個則是腦部腫瘤資料的挑戰 — BRATS。

Table 1. 兩個資料集的標籤。0為背景；1為壞死/掌骨；2為水腫/近端指骨；3為顯影腫瘤/中段指骨；4為非顯影腫瘤/遠端指骨

個人認為腦腫瘤的判讀其實不容易。(終於有個我專業的主題可以發揮一下)高惡性度的神經膠質瘤(high grade glioma/glioblastoma multiforme)，在MRI即使是人為判讀，也至少要看過兩三種造影才會比較準。磁振造影通常標準一定有T1/T2 weighted，但在腦腫瘤卻不只看這兩種，有時還要搭配一些訊號分析才能更準確的判定其邊界。由於它常常有區域性壞死(central necrosis)及周邊水腫(perifocal edema)，即使以人眼判斷邊界都可能有誤差，何況機器。

Menze[10]試著用dice coefficient的方式讓同樣的資料及給不同的人類專家所描出的邊界有個比較，以作為機器學習的參考。他們根據三個區域來評斷：
1. whole：區分正常和病變的相素
2. core：區分腫瘤核心區(即壞死和顯影/非顯影腫瘤)和周邊區(水腫及無病變區)
3. active：區分顯影腫瘤和其他。
若是單純比較兩位不同的人類專家所描出的邊界，其三項的dice coefficient分別是85%、75%和74%。若是比較一位人類專家描出的和一群專家各自描出再經過融合之後的邊界(文中稱為consensus)則提升至91%、86%和85%。因此consensus的方法所產出的邊界較有一致性和穩定性，適於作為神經網路學習的標籤。

講了這麼多，我們還是回頭來看本篇架構中所提到的一些改變在實驗中的結果吧。再整理一次相對於2D Unet，3D Unet做的調整：
1. 試著拿掉skip connection
2. 將不同大小的segmentation maps疊加在一起
3. Element-wise summation
4. 在BRATS的資料集先後測試只有一組影像、用複數組影像以及找出最好的組合
前三項是在用手部影像的資料集中完成，第四項則在BRATS的資料集完成。
不想看圖表的可以跳到下面整理作者的結論。

Figure 3. 有無skip connection對結果的影響

Figure 4. 損失函數使用Jaccard Distance和category cross-entrophy的比較

Figure 5. 比較segmentation maps有(with)無(without)疊加的結果

Figure 6. Summation和concatenation在特徵圖上的差異。concatenation可以在早期把清楚的segmentation maps呈現出來，而summation則可得到許多不相關的細節。

Figure 7. summation vs. concatenation。從統計圖表來看是難分軒輊。

Table 2. 在BRATS資料集中利用各種MRI影像訓練得到的dice score。最下面三組是組合。

Table 3. & 4. 比較BRATS系列各資料集中本架構的表現。

作者結論：
1. skip connection仍必須保留。
2. segmentation maps疊加有意義。
3. summation是略遜於concatenation的。
4. 組合複數MRI影像的訓練表現還是比較好。
5. 損失函數中雖然看起來Jaccard Distance表現比較好，但應該要將category cross entrophy做weight來當對照會更客觀(比較符合前人的另一種解決方法)

Isensee提出的概念

Figure 8. Isensee et al.提出的3D Unet模型

基本上非常相似，唯一不同似乎是左上角的input之後第一層也做了一個residual block。由於文章中是採分段敘述，以人工閱讀比較前一篇文章看到的架構差異只有這個。另外啟動函數改用Leaky ReLu。

這篇文章的另一個賣點是加上random forest試著從影像去連接預後。不過因為這次的重點擺在segmentation和Unet，這裡先略過。

Table 5. BRATS 2015資料集表現比較。模型結構相似，結果也差不多。

2D和3D：要用哪一個?

所以到底要用2D還是3D架構? 施主，這個問題應該要問你自己。

2D架構在平面或者第三維資訊較少/稀疏有較好的優勢，除了記憶體占用少，也因為可以集中精力處理平面的圖片，所以表現很突出。但若需要辨識的圖片有豐富的立體架構，甚至第三維的資訊(上下/前後)攸關辨識的結果，那選3D架構是無庸置疑的好。

有人拿了一組立體資料以2D和3D分別訓練，想看結果可以連過去瞧瞧，只能說差異驚人呀。連作者也預料不到。

以上是3D架構的介紹及2D與3D架構的比較。下一篇起將分享3D架構的實作，當然，如我之前所說，我的程式語言能力還不足以支撐自己手刻，即使有keras也還是要花很多時間心力的。所以我是找到前人的code來分析，希望能講解清楚，這樣對於後面自己的問題，可能有更多高手可以為我解答。

Reference

[1] https://paperswithcode.com/task/medical-image-segmentation
[2] B. Kayalıbay, G. Jensen, P. Smagt: CNN-based Segmentation of Medical Imaging Data. arXiv:1701.03056v2. July 2017
[3] F. Isensee, P. Kickingereder, W. Wick, M. Bendszus, K.H. Maier-Hein: Brain Tumor Segmentation and Radiomics Survival Prediction: Contribution to the BRATS 2017 Challenge. arXiv:1802.10508v1. Feb. 2018
[4] https://www.zhihu.com/question/271497684
[5] https://blog.csdn.net/MrCharles/article/details/103883043 (鎖住了)
[6] https://towardsdatascience.com/2d-or-3d-a-simple-comparison-of-convolutional-neural-networks-for-automatic-segmentation-of-625308f52aa7
[7] H. Noh, S. Hong, and B. Han. Learning deconvolution network for semantic segmentation. 2015, arXiv:1505.04366v1
[8] O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. 2015, arXiv:1505.04597v1
[9] O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. 2015, arXiv:1505.04597v1
[10] B. Menze, A. Jakab, S. Bauer, et al. The Multimodal Brain Tumor Image Segmentation Benchmark(BRATS). IEEE Transactions on Medical Imaging, page 33, 2014.