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

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

分割的準(zhǔn)確性對(duì)于醫(yī)學(xué)圖像至關(guān)重要，因?yàn)檫吘壏指铄e(cuò)誤會(huì)導(dǎo)致不可靠的結(jié)果，從而被拒絕用于臨床中。

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

UNet++里有什么新東西？

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

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

1. 重新設(shè)計(jì)的跳躍路徑（顯示為綠色）
2. 密集跳躍連接（顯示為藍(lán)色）
3. 深度監(jiān)督（顯示為紅色）

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

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

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

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

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

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

密集跳躍連接

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

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

深度監(jiān)督

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

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

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

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

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

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

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

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

在UNet++上的實(shí)驗(yàn)

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

度量

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

二元交叉熵

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

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

Dice 系數(shù)

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

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

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

IoU

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

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

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

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

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

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

• 二元交叉熵：0.2650
• Dice系數(shù)：0.8104
• IoU：0.8580

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

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


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

總結(jié)

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