羅兆林,宋亞男,徐榮華,蕭飛鵬

(廣東工業(yè)大學(xué),廣東 廣州 510006)

0 引言

醫(yī)學(xué)圖像分割主要目標(biāo)是從醫(yī)學(xué)圖像中準(zhǔn)確地識別和分離出感興趣的結(jié)構(gòu)或區(qū)域,如器官、腫瘤、血管網(wǎng)絡(luò)等。在手術(shù)過程中,準(zhǔn)確的分割可以提供實(shí)時(shí)的導(dǎo)航和可視化引導(dǎo),讓手術(shù)更加精準(zhǔn)和安全[1]。近年來,隨著計(jì)算機(jī)視覺技術(shù)的蓬勃發(fā)展,深度卷積神經(jīng)網(wǎng)絡(luò)越來越多地應(yīng)用在醫(yī)學(xué)圖像分割領(lǐng)域,分割模型的性能也在不斷提高。自從U-Net[2]網(wǎng)絡(luò)的出現(xiàn),醫(yī)學(xué)圖像分割的網(wǎng)絡(luò)結(jié)構(gòu)發(fā)生顯著的變化,從之前的單一分支的卷積神經(jīng)網(wǎng)絡(luò),發(fā)展為U型結(jié)構(gòu)的網(wǎng)絡(luò)。此后對U型結(jié)構(gòu)網(wǎng)絡(luò)的改進(jìn)噴涌而出,比如重新設(shè)計(jì)了跳躍連接的UNet++[3]、使用殘差卷積塊的R2U-Net[4]、引入了注意力機(jī)制的Attention U-Net[5]等。這些對U型網(wǎng)絡(luò)的改進(jìn)帶來了更高的性能,但它們比以往更需要大量帶標(biāo)簽數(shù)據(jù)去訓(xùn)練[6]。

影像醫(yī)生需要具備深厚的醫(yī)學(xué)知識,特別在計(jì)算機(jī)斷層掃描技術(shù)(Computed Tomography,CT)、核磁共振圖像(Magnetic Resonance Imaging,MRI)等三維醫(yī)學(xué)圖像上,復(fù)雜的三維結(jié)構(gòu)讓影像醫(yī)生標(biāo)注病灶非常耗時(shí),還有患者數(shù)據(jù)隱私等倫理問題。這些讓醫(yī)學(xué)圖像具有數(shù)量小、標(biāo)注成本高的特點(diǎn)。針對小樣本特性的醫(yī)學(xué)圖像分割研究越發(fā)重要[7]。

小樣本學(xué)習(xí)早期主要應(yīng)用于圖像分類任務(wù),其利用模型學(xué)習(xí)到的先驗(yàn)知識,輔助預(yù)測只有少量樣本的新類別。隨著在小樣本學(xué)習(xí)上的研究發(fā)展,小樣本學(xué)習(xí)已經(jīng)用于計(jì)算機(jī)視覺的機(jī)器學(xué)習(xí)的方方面面。當(dāng)前,主流的小樣本學(xué)習(xí)方法主要包括元學(xué)習(xí)(Meta-learning)、數(shù)據(jù)增強(qiáng)、遷移學(xué)習(xí)、度量學(xué)習(xí)等方法[8]。模型無關(guān)元學(xué)習(xí)(Model-agnostic Meta-learning,MAML)[9],是一種著名的元學(xué)習(xí)方法。它的目標(biāo)是讓模型學(xué)習(xí)到一個(gè)適應(yīng)性強(qiáng)的初始化參數(shù),面對新任務(wù)時(shí),讓模型經(jīng)過少量數(shù)據(jù)微調(diào)參數(shù),即可在該任務(wù)下獲得較好的性能。MAML與所使用的模型無關(guān),其能用在所有梯度下降法訓(xùn)練的模型,MAML算法應(yīng)用在圖像分類任務(wù)上非常廣泛。

目前,利用MAML元學(xué)習(xí)算法在小樣本醫(yī)學(xué)圖像分割上的研究較少。因此,本文提出了一種小樣本醫(yī)學(xué)圖像分割算法,在3D U-Net[10]網(wǎng)絡(luò)的基礎(chǔ)上做出如下改進(jìn):(1)對3D U-Net的下采樣模塊進(jìn)行改進(jìn),把每個(gè)3D卷積層的批歸一化改進(jìn)為組歸一化,增加多一層3D卷積層,從而增加了網(wǎng)絡(luò)深度與參數(shù)量,有利于結(jié)合MAML元學(xué)習(xí)算法,得到更好的調(diào)優(yōu)初始化模型,其次組歸一化能優(yōu)化網(wǎng)絡(luò)在小批量訓(xùn)練下的精度。(2)引入了Transferomer模塊,豐富網(wǎng)絡(luò)提取的全局信息。(3)在跳躍連接中引入改進(jìn)的注意力門模塊,增加組歸一化有助于確保輸入分布大致相似,Sigmoid激活函數(shù)更換為ReLU函數(shù),減少計(jì)算量以及改善梯度消失問題。(4)利用MAML元學(xué)習(xí)算法訓(xùn)練模型,增強(qiáng)網(wǎng)絡(luò)在小樣本醫(yī)學(xué)圖像數(shù)據(jù)下的表現(xiàn)。最后,與現(xiàn)有的多個(gè)方法進(jìn)行了實(shí)驗(yàn)對比分析,對本文算法的有效性進(jìn)行驗(yàn)證。

1 總體設(shè)計(jì)

1.1 模型整體結(jié)構(gòu)

圖1 整體網(wǎng)絡(luò)結(jié)構(gòu)

1.2 下采樣模塊的改進(jìn)

3D U-Net的下采樣模塊使用了2個(gè)3D卷積層串聯(lián),每個(gè)卷積層有一個(gè)3×3×3卷積核、批歸一化層以及ReLU激活函數(shù),其結(jié)構(gòu)如圖2(a)所示,輸入Fin經(jīng)過兩層3D卷積層后得到Fout。由于MAML元學(xué)習(xí)算法是一個(gè)訓(xùn)練較優(yōu)網(wǎng)絡(luò)初始化參數(shù)的過程,適當(dāng)?shù)丶哟缶W(wǎng)絡(luò)參數(shù),有利于提高網(wǎng)絡(luò)應(yīng)用在新的小樣本分割任務(wù)中的適應(yīng)能力。文本的下采樣模塊使用了3個(gè)3D卷積層串聯(lián),每個(gè)卷積層有1個(gè)3×3×3卷積核、組歸一化層以及ReLU激活函數(shù),上采樣網(wǎng)絡(luò)也是如此,其結(jié)構(gòu)如圖2(b)所示。

圖2 采樣模塊

批歸一化是深度學(xué)習(xí)中非常有效的一個(gè)技術(shù),在各種先進(jìn)的卷積神經(jīng)網(wǎng)絡(luò)都能見到其身影。批歸一化主要在批(batch)這個(gè)維度上進(jìn)行歸一化,其需要用到足夠大的批大小。3D U-Net中使用的歸一化也是批歸一化,但由于本文網(wǎng)絡(luò)增加了更多的3×3×3卷積核,加上Transformer模塊的引入,參數(shù)量增多,增加了模型訓(xùn)練難度,對顯卡顯存的需求加大。在顯存有限的條件下,只能使用更小的批次訓(xùn)練,而批量越小,批歸一化的誤差會快速增大。故本文將下采樣模塊以及上采樣模塊中的批歸一化更改為組歸一化。組歸一化把通道分為組,計(jì)算每一組之內(nèi)的均值和方差,以進(jìn)行歸一化。組歸一化的計(jì)算與批量大小無關(guān),其精度也在各種批量大小下保持穩(wěn)定。

1.3 Transformer模塊

1.4 注意力門模塊的改進(jìn)

Attention U-Net的注意力門有2個(gè)輸入,一個(gè)是來自淺層網(wǎng)絡(luò)跳躍連接的特征圖x,另一個(gè)輸入來自網(wǎng)絡(luò)每層上采樣的輸出g,經(jīng)過注意力門后得到一個(gè)尺寸與x相同的輸出x′。2個(gè)輸入分別經(jīng)過一個(gè)1×1×1的卷積核后,直接相加后輸入ReLU激活函數(shù)。本文在1×1×1卷積核后面加入了一個(gè)組歸一化層,通過標(biāo)準(zhǔn)化每個(gè)批次的輸入數(shù)據(jù),有助于確保輸入分布大致相似,從而使每個(gè)注意力門模塊都更容易學(xué)習(xí),有助于減輕梯度消失和梯度爆炸問題,提高訓(xùn)練的穩(wěn)定性。

Attention U-Net的注意力門使用了Sigmoid激活函數(shù)對注意力系數(shù)進(jìn)行歸一化處理。Sigmoid激活函數(shù)具有計(jì)算復(fù)雜、梯度消失和輸出飽和的問題,而使用Softmax激活函數(shù)會使輸出變的稀疏。本文改進(jìn)為計(jì)算更簡單、沒有梯度消失問題的ReLU函數(shù)。改進(jìn)后的注意力門結(jié)構(gòu)如圖3所示。

圖3 改進(jìn)的注意力門結(jié)構(gòu)

2 實(shí)驗(yàn)與結(jié)果分析

2.1 實(shí)驗(yàn)數(shù)據(jù)及預(yù)處理

本文實(shí)驗(yàn)使用多中心、多供應(yīng)商和多疾病心臟圖像分割挑戰(zhàn)賽(Multi-Centre, Multi-Vendor &Multi-Disease Cardiac Image Segmentation Challenge,M &Ms)[12]數(shù)據(jù)集,該數(shù)據(jù)集包含320份樣本,這些樣本使用4臺不同的核磁共振儀器采集固有ABCD 4個(gè)不同的域的樣本。其中,域A有95份樣本,域B有125份樣本,域C和域D各包含50份樣本。該數(shù)據(jù)集提供了3個(gè)真值標(biāo)簽,分別為左心室(Left Ventricle,LV)、右心室(Right Ventricle,RV)和左心室心肌(Myocardium,MYO)。

本文的數(shù)據(jù)集劃分如表1所示,域A與域B的數(shù)據(jù)用于支持集Ds,域C的數(shù)據(jù)用于驗(yàn)證集Dv,域D的數(shù)據(jù)用于測試集Dt。

表1 數(shù)據(jù)集設(shè)置

2.2 損失函數(shù)

交叉熵?fù)p失是圖像分割任務(wù)的常用損失函數(shù)之一。交叉熵?fù)p失會計(jì)算圖像的每一個(gè)像素的類預(yù)測,然后取平均值。但部分圖像分割任務(wù),例如:醫(yī)學(xué)圖像心臟分割,心臟只占了胸腔很小一部分,即圖像真值標(biāo)簽只占圖像很小一部分,背景占據(jù)了更大一部分。在交叉熵?fù)p失下,模型學(xué)習(xí)到了更多的背景類像素,而只學(xué)習(xí)到很小一部分真值類像素。

Dice損失在醫(yī)學(xué)圖像分割任務(wù)中比交叉熵?fù)p失更為常見。其計(jì)算所有像素的真值標(biāo)簽與預(yù)測標(biāo)簽的交并比。計(jì)算交并比就不會引入大量無關(guān)的背景像素,可以極大減緩真值與背景類別不平衡的問題。

考慮上述2種損失函數(shù)的優(yōu)劣,本文使用交叉熵?fù)p失和Dice損失的未加權(quán)和,其式如式(1)。

L=LCE+LDice

(1)

LCE代表交叉熵?fù)p失,LDice代表Dice損失。

這樣的復(fù)合損失函數(shù)被證明在各種分割任務(wù)中具有魯棒性[13]。

2.3 實(shí)驗(yàn)環(huán)境與參數(shù)設(shè)置

實(shí)驗(yàn)環(huán)境:中央處理器(I9-10900K @3.7 GHz,Intel,美國),獨(dú)立顯卡(GeForce 3090 24 GB,Nvidia,美國)。深度學(xué)習(xí)框架為PyTorch 1.11.1(Linux Foundation,美國),編程語言為Python 3.8.6(Python Software Foundation,美國)。內(nèi)循環(huán)使用學(xué)習(xí)率為0.5的普通梯度下降方法,外循環(huán)為學(xué)習(xí)率0.0001、權(quán)重衰減1e-5的Adam優(yōu)化器。

2.4 評價(jià)指標(biāo)

本文性能度量指標(biāo)選取戴斯相似性系數(shù)(Dice Similarity Coefficient,DSC)以及豪斯多夫距離(Hausdorff Distance,HD)評估每個(gè)模型的分割精度。DSC可以衡量手動(dòng)注釋和預(yù)測結(jié)果之間的重疊程度。HD可以評價(jià)模型預(yù)測結(jié)果與手動(dòng)注釋的形狀相似度,其值越小越相似,最小值為0。為了評估模型在2個(gè)指標(biāo)下的綜合表現(xiàn),本文提出了融合指標(biāo)COMBINE,其公式如式(3)。

(2)

該指標(biāo)為DSC的一半與(100-HD)一半的和,DSC越大,HD越小,融合指標(biāo)越高。反之,DSC越小,HD越大,融合指標(biāo)越低。

2.5 實(shí)驗(yàn)結(jié)果與分析

表2測試了主流醫(yī)學(xué)圖像分割網(wǎng)絡(luò)3D U-Net、UNETR[14]、RegUNet[15]、TransBTS以及Attention U-Net。實(shí)驗(yàn)結(jié)果表明,本文算法在不使用MAML算法下,對比主流分割網(wǎng)絡(luò),有更好的性能。MAML+3D U-Net網(wǎng)絡(luò)比基線3D U-Net網(wǎng)絡(luò),DSC分?jǐn)?shù)提高的同時(shí),HD分?jǐn)?shù)也下降了,表明網(wǎng)絡(luò)通過MAML算法訓(xùn)練后,分割結(jié)果不但與真值重合度高,而且分割結(jié)果的形狀與真值更接近,MAML算法帶來的提升在Attention U-Net以及本文算法上,均有所體現(xiàn)。本文算法結(jié)合MAML算法后,得到了測試中最高的平均DSC得分以及最低的平均HD分?jǐn)?shù)。

表2 實(shí)驗(yàn)評估結(jié)果

圖4展示了各個(gè)方法在測試集某個(gè)樣本的分割結(jié)果可視化,圖4右下角真值圖片標(biāo)注了右心室(RV)、左心室(LV)、左心室心肌(MYO)對應(yīng)的真值掩碼。UNETR在測試集下分割效果較差,結(jié)合表2該網(wǎng)絡(luò)的數(shù)據(jù),表明該網(wǎng)絡(luò)在測試集與訓(xùn)練集域不一致環(huán)境下,適應(yīng)性較差。在測試樣本下,Attention U-Net的分割效果會比本文的模型稍好,右心室區(qū)域過度分割的更少。在結(jié)合MAML元學(xué)習(xí)算法之后,本文的網(wǎng)絡(luò)分割效果最為接近真值標(biāo)簽。圖4縱向?qū)Ρ?個(gè)網(wǎng)絡(luò)有無使用MAML算法的結(jié)果,能看到MAML算法并非都能起到作用,對于3D U-Net網(wǎng)絡(luò),使用MAML算法后,右心室區(qū)域的過度分割更加嚴(yán)重。對于本文的網(wǎng)絡(luò),MAML算法起到了良好的效果,改善了原本右心室區(qū)域的過度分割。

圖4 分割結(jié)果可視化對比

2.6 消融實(shí)驗(yàn)

為了驗(yàn)證引入Transformer模塊以及注意力門改進(jìn)帶來的效果,本文進(jìn)行了各模塊的消融實(shí)驗(yàn)。各個(gè)模塊消融實(shí)驗(yàn)?zāi)Ｐ褪褂昧讼嗤臄?shù)據(jù)集和參數(shù)訓(xùn)練,3D U-Net+Transformer是單獨(dú)在3D U-Net引入Transformer模塊,3D U-Net+改進(jìn)注意力門是單獨(dú)引入改進(jìn)注意力門模塊。實(shí)驗(yàn)結(jié)果如表3所示。

實(shí)驗(yàn)結(jié)果可以看出:

(1)在引入Transformer模塊后的3D U-Net網(wǎng)絡(luò),能取得比基線更好的分割效果,DSC系數(shù)得到了

提升的同時(shí),HD也有所下降,證明Transformer在數(shù)據(jù)量有限的小樣本任務(wù)下,也能給模型帶來精度的提升。

(2)在單獨(dú)引入注意力門模塊后,模型性能對比3D U-Net基線得到了提升,注意力門能幫助模型消除來自淺層網(wǎng)絡(luò)的噪聲,提高對任務(wù)相關(guān)信息的專注度。

3 結(jié)語

本文在3D U-Net網(wǎng)絡(luò)的基礎(chǔ)上,對其下采樣模塊增多了一層3D卷積層,將其中的批歸一化替換為了組歸一化;將該網(wǎng)絡(luò)編解碼器連接處替換為Transformer,又在跳躍連接處,引入了改進(jìn)的注意力門模塊;利用MAML元學(xué)習(xí)算法訓(xùn)練網(wǎng)絡(luò),提升網(wǎng)絡(luò)在小樣本任務(wù)上的泛化能力。在公開數(shù)據(jù)集M&Ms上的實(shí)驗(yàn)表明,與多個(gè)現(xiàn)有方法相比,本文的算法表現(xiàn)出更好的分割性能。但本文方法增加了卷積核、Transformer以及注意力門模塊,增加了算法的復(fù)雜度以及推理速度,下一步探索輕量化算法,實(shí)現(xiàn)分割精度和算法復(fù)雜度平衡,在其他醫(yī)學(xué)圖像數(shù)據(jù)集上進(jìn)行推廣改進(jìn),提高方法的普適性。