Unet小簡介 — 實作2D Unet

最近大半年都幾乎沒有更新，主要是在忙著一些研究的事情。除此之外，我們也開始將研究計畫的內容申請專利，所以額外要處理專利申請書的事項。

今天趁著過年的小空檔，把之前自己用pytorch實作的Unet結構和大家分享。關於2D Unet的架構介紹，歡迎前往我之前寫的這篇閱覽。

程式碼是從這邊改的：https://github.com/milesial/Pytorch-UNet
我最後會說明更改的地方。

從convolutional block開始

由於encoding和decoding都使用卷積，所以可以從這個block物件著手。它是Unet架構裡的磚頭。

我選用三層卷積組成這個block。依據你的需求，可以增減卷積層數。每一層的卷積後面都搭配batch normalization以降低梯度消失的機會。在輸出卷積block之前，經過一次Relu啟動函數。至於池化層為什麼沒做進來?因為那是只有encoding才會用到的，如果在這邊放進來，到了decoding就不能用了。

再來是關於pytorch的語法。pytorch允許調用它建立的各種方法，建立在python裡面的類別繼承方式(super)之上。所以是繼承了pytorch裡面的對應物件類別之後，再根據自己的需求建立屬性，然後調用。

在建立卷積層的時候，和keras不同，必須要自己計算清楚每個輸入和輸出之後的feature map大小，這樣串接的時候才不會跳錯。有更多設計彈性的代價就是必須把細節指定更清楚。

sequential則是pytorch建立模型時最常用的方法之一，直接告訴pytorch按照你要求的順序逐層把模型拼起來。最後，要用forward方法告訴pytorch正向傳播的路線；反向傳播時pytorch會自己反過來操作。

看一下程式碼吧。

建立Unet架構

磚塊做好了，接著就可以疊模型了。要考慮的有以下幾點：
1. 輸入：輸入是幾個channel。
2. encoding(down sampling)：下降段要幾層。這會關係到上升段就需要幾層。
3. 底層：有些做法是把底層視為下降段的一部分，我則是把它拿出來單獨設計，原因是因為我設計的單層下降段裡面包含池化層，但底層不需要再池化。
4. 輸出：在上升段的最頂層經過一個CNN再輸出。

首先先設計一層下降段的模型。一層下降段應該包含一個卷積block，以及一個池化層。我選擇maxpooling作為池化層，你也可以用其他的pooling方法。

再來設計一層上升段的模型。一層上升段包含來自前一層的upsample，然後再進入卷積block。先upsample的原因是，從底層出來之後，是先經過upsample的。至於concatenate，不用在這邊處理，等一下在後面forward的時候再做。

好了，這下樓面也完成，直接把樓層疊起來吧。這邊我使用for loop的方式，算是一個和原版最大的差異：我希望我建立模型的時候不需要受到層數的限制，愛疊幾層都可以。為了這個目的，用for loop就不需要一層層限制跟明訂，可以保有彈性。

最後講一下forward裡面的做法。由於在upsample之後得到的feature map大小，有可能和同一層要進入下個downsample之前的feature map有一點誤差(誤差主要是在feature map的長寬不是偶數時發生)，所以會用padding的方式來彌補差距。彌補完之後再用concatenation，然後就可以進入下一層上升段了。

用個小函數幫助建立實例(instance)，要調用時就直接呼叫函數即可。

來看一下這一段的程式吧。

整個實作裡面，我覺得最容易出錯的地方是feature map的維度。只要沒算好就會跳錯，因此花了不少時間在檢查上面。

以上分享給大家。之後如果還有什麼新的概念或想法，我還是會持續更新的XD祝大家新年快樂。