UNet++與Ronneberger的U-Net進(jìn)行比較

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

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

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

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

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

UNet++里有什么新東西？

下面是UNet++和U-Net架構(gòu)的示意圖。

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

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

密集跳躍連接（顯示為藍(lán)色）

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

重新設(shè)計的跳躍路徑

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

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

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

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

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

密集跳躍連接

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

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

深度監(jiān)督

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

對于accurate模式，所有分割分支的輸出進(jìn)行平均。

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

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

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

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

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

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

在UNet++上的實(shí)驗

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

度量

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

二元交叉熵

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

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

Dice 系數(shù)

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

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

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

IoU

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

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

訓(xùn)練和結(jié)果

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

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

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

二元交叉熵：0.2650

Dice系數(shù)：0.8104

IoU：0.8580

幾個U-Net模型之間的指標(biāo)進(jìn)行比較，如下所示。

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


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

總結(jié)

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