【論文筆記】基於U-Net的醫學圖像分割模型nnU-Net（no-new U-Net）

本文是 nnU-Net 論文的筆記，稍微做了整合了一下，可能有些地方有所遺漏，等以後看完代碼之後再補充。

1. 摘要

nnU-Net（no-new-Net）是基於 U-Net 的一種模型，它在達到很好的效果的同時還實現了自適應（self-adapting）的功能。在不同數據集（或不同的部位）的醫學圖像上進行分割時，往往需要具有不同結構的網絡和不同的訓練方案，自適應是指模型在對不同的數據集進行訓練時，可以自動的調整 batch-size、patch-size 等，以達到很好的效果。

作者在一個醫學圖像分割十項全能比賽當中的 6 個數據集集中都取到了當時最好的結果，這個比賽一共 10 個數據集，它會給出 7 個數據集讓你訓練，然後在其他 3 個數據集上進行驗證。

2. 方法

2.1 網絡結構

作者以 U-Net 爲基礎，通過調整網絡的結構後效果沒有得到提升，所以作者轉而關注非結構的部分。nnU-Net 可以根據給出的數據集自動的調整網絡的設置。nnU-Net 相比於 U-Net，前者用 leaky ReLU 代替了 ReLU 做激活函數，用實例正則化（instance normalization）代替了批正則化（batch normalization）。

U-Net 採用編碼器-解碼器的結構，其中編碼器部分可以聚合語義信息（semantic information），但是同時也會損失空間信息（spatial information），說人話就是編碼器在提取圖像特徵的同時也使圖像的分辨率降低了。但是語義信息和空間信息對於分割任務來說都是至關重要的，所以在解碼器部分需要對空間信息進行恢復，具體的就是通過下層的語義信息，並結合利用跳躍連接得到的更高分辨率的特徵圖（即空間信息）來進行空間信息的恢復。

nnU-Net 由三種基礎的 U-Net 網絡組成，分別是 2D U-Net，3D U-Net 和 U-Net Cascade。其中，2D 和 3D U-Net 產生一個全像素的分割結果，U-Net Cascade 先產生一個低像素的分割結果，再對其進行微調。

直觀來看，2D U-Net 在處理 3D 的醫學圖像時不是最優的，因爲在 $z$ 軸上的信息沒有得到充分的利用。但是也有證明顯示 3D 分割方法在一些具有各向異性（不懂何爲各向異性）的數據集上效果會變差。我們通常會選擇 3D U-Net 來處理 3D 的醫學圖像，但這 3D 分割方法的一個缺點就是佔用顯存較大，當顯存放不開一個 3D 圖像時，可以將圖像的 patch 作爲輸入，但這又限制了視野（field of view），不能充分的利用圖像的上下文信息。爲了解決 3D U-Net 帶來的這種問題，所以使用了 U-Net Cascade 的方法，即先用一個 3D U-Net 處理下采樣之後的圖像（圖像分辨率較低），然後輸出分割結果，並將其上採樣到原始大小，並作爲附加的輸入通道（和原始 3D 圖像在通道上拼接）輸入到第二個 3D U-Net 做基於 patch 的分割。

上圖是 U-Net Cascade 的示意圖，要注意每個不同標記所代表的意思，然後結合上一段的內容還是比較好理解的。

2.2 動態自適應

網絡的動態自適應簡單來說就是根據數據集的不同自動的調整 batch size 和 patch size，當 patch size 大的時候 batch size 就小一點；當前者小的時候，後者就大一點。此外還要調整池化和卷積操作的個數。

3. 預處理

3.1 裁剪

將圖像中非零的區域裁掉，可以減小顯存的負擔。

3.2 重採樣

由於成像設備和成像協議的不同，醫學圖像的體像素間距也不同，所以要對其進行重採樣。具體的，對醫學圖像使用三階樣條插值法，對其相應的分割標籤使用最近鄰插值法。

3.3 正則化

對於 CT 圖像而言，對其像素值歸一化到原像素值的 0.5%~95.5%，並根據均值和方差對其進行 z-score 正則化。對於非 CT 圖像而言，只進行 z-score 正則化。

4. 訓練

4.1 網絡設置

訓練時的總損失 $= dice$ 損失 + 交叉熵損失，優化器使用的是 Adam 優化器，學習率爲 $3\times 10^{-4}$，每當訓練損失的指數移動平均值在最近 30 個 epoch 內沒有改善時，學習率下降 0.2 倍。當在 60 個 epoch 內，學習率的變化不超過 $5\times10^{-3}$ 並且學習率不小於 $10^{-6}$ 時，停止訓練。在驗證時使用五折交叉驗證。

4.2 數據增強

使用隨機旋轉、隨機縮放、隨機彈性變形、伽瑪校正增強和反射做數據增強。在第二階段的 U-Net Cascade 中使用了形態學操作和隨機移除分割中的連接組件來進行數據增強。

5. 推斷

因爲網絡的準確率在一個 patch 中的中央部分較高，而在邊緣部分較低，所以在在中央部分的體素的權重較高，而邊緣部分的較低。

6. 後處理

對 3D 醫學圖像進行分割，最終器官的分割結果在空間當中一定是連通的，所以最後只保存最大的連通分量，其他的均刪除。