一種改進的小樣本醫學圖像分割算法研究

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

(廣東工業大學,廣東 廣州 510006)

0 引言

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

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

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

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

1 總體設計

1.1 模型整體結構

圖1 整體網絡結構

1.2 下采樣模塊的改進

3D U-Net的下采樣模塊使用了2個3D卷積層串聯,每個卷積層有一個3×3×3卷積核、批歸一化層以及ReLU激活函數,其結構如圖2(a)所示,輸入Fin經過兩層3D卷積層后得到Fout。由于MAML元學習算法是一個訓練較優網絡初始化參數的過程,適當地加大網絡參數,有利于提高網絡應用在新的小樣本分割任務中的適應能力。文本的下采樣模塊使用了3個3D卷積層串聯,每個卷積層有1個3×3×3卷積核、組歸一化層以及ReLU激活函數,上采樣網絡也是如此,其結構如圖2(b)所示。

圖2 采樣模塊

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

1.3 Transformer模塊

1.4 注意力門模塊的改進

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

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

圖3 改進的注意力門結構

2 實驗與結果分析

2.1 實驗數據及預處理

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

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

表1 數據集設置

2.2 損失函數

交叉熵損失是圖像分割任務的常用損失函數之一。交叉熵損失會計算圖像的每一個像素的類預測,然后取平均值。但部分圖像分割任務,例如:醫學圖像心臟分割,心臟只占了胸腔很小一部分,即圖像真值標簽只占圖像很小一部分,背景占據了更大一部分。在交叉熵損失下,模型學習到了更多的背景類像素,而只學習到很小一部分真值類像素。

Dice損失在醫學圖像分割任務中比交叉熵損失更為常見。其計算所有像素的真值標簽與預測標簽的交并比。計算交并比就不會引入大量無關的背景像素,可以極大減緩真值與背景類別不平衡的問題。

考慮上述2種損失函數的優劣,本文使用交叉熵損失和Dice損失的未加權和,其式如式(1)。

L=LCE+LDice

(1)

LCE代表交叉熵損失,LDice代表Dice損失。

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

2.3 實驗環境與參數設置

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

2.4 評價指標

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

(2)

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

2.5 實驗結果與分析

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

表2 實驗評估結果

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

圖4 分割結果可視化對比

2.6 消融實驗

為了驗證引入Transformer模塊以及注意力門改進帶來的效果,本文進行了各模塊的消融實驗。各個模塊消融實驗模型使用了相同的數據集和參數訓練,3D U-Net+Transformer是單獨在3D U-Net引入Transformer模塊,3D U-Net+改進注意力門是單獨引入改進注意力門模塊。實驗結果如表3所示。

實驗結果可以看出:

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

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

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

3 結語

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