UNet++與Ronneberger的U-Net進行比較

使用一系列的網(wǎng)格狀的密集跳躍路徑來提升分割的準確性。

在這篇文章中，我們將探索UNet++: A Nested U-Net Architecture for Medical Image Segmentation這篇文章，作者是亞利桑那州立大學的Zhou等人。本文是U-Net的延續(xù)，我們將把UNet++與Ronneberger等人的U-Net原始文章進行比較。

UNet++的目標是通過在編碼器和解碼器之間加入Dense block和卷積層來提高分割精度。

分割的準確性對于醫(yī)學圖像至關重要，因為邊緣分割錯誤會導致不可靠的結果，從而被拒絕用于臨床中。

為醫(yī)學成像設計的算法必須在數(shù)據(jù)樣本較少的情況下實現(xiàn)高性能和準確性。獲取這些樣本圖像來訓練模型可能是一個消耗資源的過程，因為需要由專業(yè)人員審查的高質量、未壓縮和精確注釋的圖像。

UNet++里有什么新東西？

下面是UNet++和U-Net架構的示意圖。

UNet++ 在原始的U-Net上加了3個東西：

重新設計的跳躍路徑（顯示為綠色）

密集跳躍連接（顯示為藍色）

深度監(jiān)督（顯示為紅色）

重新設計的跳躍路徑

在UNet++中，增加了重新設計的跳躍路徑(綠色顯示)，以彌補編碼器和解碼器子路徑之間的語義差別。

這些卷積層的目的是減少編碼器和解碼器子網(wǎng)絡的特征映射之間的語義差距。因此，對于優(yōu)化器來說，這可能是一個更直接的優(yōu)化問題。

U-Net采用跳躍連接，直接連接編碼器和解碼器之間的特征映射，導致把語義上不相似的特征映射相融合。

然而，在UNet++中，相同dense block的前一個卷積層的輸出與較低層dense block對應的上采樣輸出進行融合。這使得已編碼特征的語義級別更接近于等待在解碼器中的特征映射的語義級別，因此，當接收到語義上相似的特征映射時，優(yōu)化更容易。

跳躍路徑上的所有卷積層使用大小為3×3的核。

密集跳躍連接

在UNet++中，密集跳躍連接(用藍色顯示)實現(xiàn)了編碼器和解碼器之間的跳躍路徑。這些Dense blocks是受到DenseNet的啟發(fā)，目的是提高分割精度和改善梯度流。

密集跳躍連接確保所有先驗特征圖都被累積，并通過每個跳躍路徑上的dense卷積塊而到達當前節(jié)點。這將在多個語義級別生成完整分辨率的特征映射。

深度監(jiān)督

UNet++中增加了深度監(jiān)督(紅色顯示)，通過修剪模型來調整模型的復雜性，在速度(推理時間)和性能之間實現(xiàn)平衡。

對于accurate模式，所有分割分支的輸出進行平均。

對于fast模式，最終的分割圖從分割分支之一選擇。

Zhou等人進行了實驗，以確定在不同剪枝水平下的最佳分割性能。使用的度量是IoU和推理時間。

他們試驗了四種分割任務：a)細胞核，b)結腸息肉，c)肝臟，d)肺結節(jié)。結果如下：

L3與L4相比，平均減少了32.2%的推理時間，同時IoU略微降低了。

更激進的修剪方法，如L1和L2，可以進一步減少推理時間，但以顯著的分割性能為代價。

在使用UNet++時，我們可以調整用例的層數(shù)。

在UNet++上的實驗

我使用Drishti-GS數(shù)據(jù)集，這與Ronneberger等人在他們的論文中使用的不同。該數(shù)據(jù)集包含101幅視網(wǎng)膜圖像，以及用于檢測青光眼的光學disc和光學cup的mask標注。青光眼是世界上致盲的主要原因之一。50張圖片用于訓練，51張用于驗證。

度量

我們需要一組指標來比較不同的模型，這里我們有二元交叉熵，Dice 系數(shù)和IoU。

二元交叉熵

二分類的常用度量和損失函數(shù)，用于度量誤分類的概率。

我們將使用PyTorch的binary_cross_entropy_with_logits函數(shù)，與Dice系數(shù)一起作為損失函數(shù)對模型進行訓練。

Dice 系數(shù)

對預測值和實際值之間重疊的通用的度量。計算方法為 2 × 預測值與ground truth的重疊面積除以預測值與ground truth之和的總面積。

這個度量指標的范圍在0到1之間，其中1表示完全重疊。

我使用這個度量和二元交叉熵作為訓練模型的損失函數(shù)。

IoU

一個簡單(但有效！)的度量來計算預測的mask與ground truth mask的準確性。計算預測值與ground truth的重疊面積，并除以預測值與ground truth的并集面積。

類似于Dice系數(shù)，這個度量指標的范圍是從0到1，其中0表示沒有重疊，而1表示預測與地面真實之間完全重疊。

訓練和結果

優(yōu)化這個模型，訓練50多個epoch，使用Adam優(yōu)化器，學習率1e-4，學習率衰減率為沒10個epochs乘以0.1， 損失函數(shù)是二元交叉熵和Dice 系數(shù)的組合。

模型在27分鐘內完成了36.6M可訓練參數(shù)的訓練，每個epoch大約需要32秒。

表現(xiàn)最好的epoch是第45個epoch(在50個epochs中)。

二元交叉熵：0.2650

Dice系數(shù)：0.8104

IoU：0.8580

幾個U-Net模型之間的指標進行比較，如下所示。

測試首先通過模型處理一些沒見過的樣本，來預測光學disc(紅色)和光學cup(黃色)。這里是UNet++和U-Net的測試結果對比。


從指標表來看，UNet++在IoU上超越U-Net，但在Dice系數(shù)上表現(xiàn)較差。從定性測試的結果來看，UNet++成功地對第一張圖像進行了正確的分割，而U-Net做得并不好。也由于UNet++的復雜性，訓練時間是U-Net的兩倍。必須根據(jù)它們的數(shù)據(jù)集評估每種方法。

總結

UNet++的目標是提高分割精度，通過一系列嵌套的，密集的跳躍路徑來實現(xiàn)。重新設計的跳躍路徑使語義相似的特征映射更容易優(yōu)化。密集的跳躍連接提高了分割精度，改善了梯度流。深度監(jiān)督允許模型復雜性調優(yōu)，以平衡速度和性能優(yōu)化。