冠軍方案解讀 | nnUNet改進(jìn)提升筆記
本文描述了基于nn-UNet試驗了幾種改進(jìn)，包括使用更大的網(wǎng)絡(luò)、用GN替換BN以及在解碼器中使用Axial Attention。與Baseline相比，量化指標(biāo)略有改進(jìn)。在unseen test data的最終排名中，nn-UNet贏得了第一名的好成績。
1改進(jìn)策略
1.1 數(shù)據(jù)方面
BraTS2021包括了來自2000名患者的多參數(shù)MRI掃描結(jié)果，其中1251人的圖像提供了分割標(biāo)簽給參與者來開發(fā)算法，其中219人在驗證階段被用于公共排行榜，其余530個案例用于私人排行榜和參與者的最終排名。
MRI掃描有4種對比:
原生T1加權(quán)圖像 對比后T1加權(quán)(T1GD) T2加權(quán) T2流體衰減反轉(zhuǎn)恢復(fù)(T2-Flair)
注釋由1-4個評分員手工完成，最終得到有經(jīng)驗的神經(jīng)放射學(xué)家的批準(zhǔn)。這些標(biāo)簽包括gd增強(qiáng)腫瘤(ET)、瘤周水腫/侵襲組織(ED)和壞死腫瘤核心(NCR)區(qū)域。所有MRI掃描均通過同解剖模板配準(zhǔn)、各向同性1mm3分辨率插值和顱骨剝離進(jìn)行預(yù)處理。所有MRI掃描及相關(guān)標(biāo)記的圖像大小為240×240×155。
圖1顯示了這4種對比與分割的代表性切片。在輸入網(wǎng)絡(luò)之前，對提供的數(shù)據(jù)進(jìn)行了進(jìn)一步的處理。為了減少計算量，將volumes裁剪為non-zero voxels。由于MR圖像的強(qiáng)度是定性的，因此根據(jù)其均值和標(biāo)準(zhǔn)差對voxels進(jìn)行歸一化。
1.2 模型方面
1、Baseline nnU-Net
nnU-Net核心是一個在128×128×128大小的Patch上運行的3D U-Net。該網(wǎng)絡(luò)具有編碼-解碼器結(jié)構(gòu)，并帶有Skip Connection，將兩條路徑連接起來。
該編碼器由5個相同分辨率的卷積層組成，具有卷積下采樣功能。該解碼器遵循相同的結(jié)構(gòu)，使用轉(zhuǎn)置卷積上采樣和卷積操作在同一級別上的編碼器分支的串聯(lián)Skip特征。每次卷積操作后，采用斜率為0.01的Leaky ReLU(lReLU)和批歸一化處理。mpMRI volumes被連接并作為4通道輸入。
nnU-Net應(yīng)用Region-Based訓(xùn)練，而不是預(yù)測3個相互排斥的腫瘤子區(qū)域，而不是預(yù)測3個互斥腫瘤分區(qū),與提供的分割標(biāo)簽一樣，該網(wǎng)絡(luò)預(yù)測的是增強(qiáng)腫瘤的3個重疊區(qū)域如加強(qiáng)腫瘤(ET,original region),腫瘤核心或TC(ET+necrotic tumor),和整個腫瘤或WT(ET+NT+ED)。
網(wǎng)絡(luò)的最后一層的softmax被sigmoid所取代，將每個voxels作為一個多類分類問題。
由于公共和私人排行榜的計算指標(biāo)是基于這些區(qū)域的，這種基于區(qū)域的訓(xùn)練可以提高表現(xiàn)。額外的sigmoid輸出添加到每個分辨率除了2個最低的水平，應(yīng)用深度監(jiān)督和改善梯度傳播到早期層。卷積濾波器的數(shù)量被初始化為32個，并且分辨率每降低一倍，最大可達(dá)320個。
2、更大的網(wǎng)絡(luò)和GN
第一個修改是，通過將編碼器中的kernel數(shù)量加倍，同時在解碼器中保持相同的kernel，非對稱地增加了網(wǎng)絡(luò)的大小。由于訓(xùn)練數(shù)據(jù)的數(shù)量是前一年的4倍，增加網(wǎng)絡(luò)的容量將有助于它能夠建模更大的數(shù)據(jù)種類。kernel的最大數(shù)量也增加到512個，改進(jìn)后的網(wǎng)絡(luò)結(jié)構(gòu)如圖2所示：
第二個修改是用GN代替所有BN。即使使用混合精度訓(xùn)練，3D卷積網(wǎng)絡(luò)也需要大量的GPU內(nèi)存，這限制了在訓(xùn)練中可以使用的Batch-Size。
3、Axial attention解碼器
最后添加的是在解碼器中使用Axial attention。Self-Attention或Transformer是一個突破性的想法，允許學(xué)習(xí)一個輸入序列的自適應(yīng)注意力僅僅基于它自己。Self-Attention最初是在自然語言處理中，現(xiàn)在已經(jīng)慢慢被計算機(jī)視覺研究所采用。當(dāng)試圖將Self-Attention應(yīng)用于視覺問題時，主要的障礙之一是注意力機(jī)制的計算復(fù)雜度與輸入的大小成二次方，這使得它不可能適合或訓(xùn)練網(wǎng)絡(luò)在一個標(biāo)準(zhǔn)的工作站設(shè)置。當(dāng)處理帶有額外維度的3D數(shù)據(jù)時，這是一個更大的問題。
Axial attention最近被提出作為將注意力應(yīng)用于多維數(shù)據(jù)時的一種有效解決方案。通過將Self-Attention獨立地應(yīng)用于輸入的每一個軸上，計算只與圖像大小成線性比例，使注意力機(jī)制即使與3D數(shù)據(jù)整合成為可能。
本文將Axial attention應(yīng)用到網(wǎng)絡(luò)的解碼器上，將其運行在轉(zhuǎn)置卷積上采樣的輸出上，然后將它們相加。
顯示了Axial attention解碼器塊的示意圖。即使有了更有效的注意力，作者發(fā)現(xiàn)這種方法也不可能應(yīng)用于最高分辨率的特征(128×128×128)，因此只選擇了4個較低分辨率的特征。注意力頭的數(shù)量和每個頭的尺寸隨著分辨率的降低而增加一倍，分別從4和16(64×64×64分辨率)開始。
1.3 訓(xùn)練策略
這里遵循nnU-Net訓(xùn)練方法。每個網(wǎng)絡(luò)都接受了5倍交叉驗證的訓(xùn)練。在訓(xùn)練過程中，動態(tài)地應(yīng)用數(shù)據(jù)增強(qiáng)來提高泛化能力。數(shù)據(jù)增強(qiáng)包括隨機(jī)旋轉(zhuǎn)和縮放、彈性變形、附加亮度增強(qiáng)和伽瑪縮放。
優(yōu)化的目標(biāo)是binary entropy loss和Dice loss的總和，計算在最終的全分辨率輸出以及在低分辨率的輔助輸出。使用Batch Dice loss代替sample Dice loss，將整個批次作為一個樣本來計算損失，而不是平均每個樣本在小批中的Dice。批量Dice幫助穩(wěn)定訓(xùn)練通過減少來自樣本的錯誤，少量的注釋樣本。網(wǎng)絡(luò)采用Nesterov動量為0.99的隨機(jī)梯度下降法進(jìn)行優(yōu)化。初始學(xué)習(xí)率為0.01，并按照多項式計劃衰減：
每次訓(xùn)練運行持續(xù)1000個epoch，每個epoch包含250個小批量。當(dāng)前折疊驗證集上的Dice Score用于監(jiān)視訓(xùn)練進(jìn)度。所有實驗都是在24GB VRAM的NVIDIA RTX 3090 GPU上使用Pytorch 1.9進(jìn)行的。開發(fā)了以下模型:
BL：Baseline nnUNet BL+L：Baseline with Large nnUNet BL+GN：Baseline with Group Normalization BL+AA：Baseline with axial attention, batch normalization BL+L+GN：nnUNet with larger Unet, group normalization
2比賽結(jié)果
3參考閱讀
[1].Extending nn-UNet for brain tumor segmentation
4推薦閱讀
全新范式 | Box-Attention同時讓2D、3D目標(biāo)檢測、實例分割輕松漲點（求新必看）
Swin-Transformer又又又下一城 | 看SwinTrack目標(biāo)跟蹤領(lǐng)域獨領(lǐng)風(fēng)騷
全面超越Swin Transformer | Facebook用ResNet思想升級MViT
長按掃描下方二維碼添加小助手。
可以一起討論遇到的問題
聲明：轉(zhuǎn)載請說明出處
掃描下方二維碼關(guān)注【集智書童】公眾號，獲取更多實踐項目源碼和論文解讀，非常期待你我的相遇，讓我們以夢為馬，砥礪前行！
