用于多模態MRI腦腫瘤圖像分割的融合雙重對抗學習CNN-Transformer模型

華楷文,方賢進

(安徽理工大學 計算機科學與工程學院,安徽 淮南 232001)

膠質瘤是人類最為常見的原發性腫瘤之一。目前膠質瘤的影像檢查一般使用計算機斷層掃描(computed tomography,CT)和核磁共振成像(magnetic resonance imaging,MRI),通過這2種方式可以準確查看到膠質瘤病變的形態以及準確位置。其中,MRI檢查在識別腦組織結構上有較高的分辨率,一般作為輔助腦部膠質瘤治療的首選手段。在臨床應用上,經驗不同的醫生對影像上病變區域進行手動標注和分割時會產生一定差異,這將給最終診斷帶來巨大的困難。因此,在MRI圖像上是否能準確分割出病變區域對于膠質瘤的輔助治療十分重要。

隨著人工智能和深度學習的快速發展,卷積神經網絡(convolutional neural network,CNN)在醫學圖像分割領域取得了極大成功[1]。在腦腫瘤圖像分割領域,Colman等[2]將殘差模塊引入U-Net模型,將其中104層卷積塊替換成殘差塊,提出具有 104 個卷積層的深度殘差 UNet(deep residual UNet with 104 convolutional layers,DR-UNet104)模型。Awasthi等[3]開發了基于注意力機制的U-Net分割模型,該模型用于在MRI影像中識別腫瘤的各個區域。Xu等[4]將注意力模塊和深度監督方法結合,提出了用于腦腫瘤圖像分割的U-Attention Net框架。?i?ek等[5]在U-Net模型基礎上提出了針對3D圖像分割的3D U-Net模型,該模型只需將U-Net中的2D卷積塊轉換成3D卷積塊。3D U-Net模型的提出使得3D醫學圖像的分割性能得到了極大提升。在此基礎上,產生了一系列3D U-Net腦腫瘤圖像分割模型。如Zhao等[6]在3D U-Net模型基礎上提出名為“Categorical Dice”的新型損失函數,通過為不同分割區域設置權重解決了前景和后景圖像體素不平衡的問題。何康輝等[7]通過將3D U-Net模型中通道數減少到原來的1/4來降低模型參數量,同時在模型中使用空間-通道壓縮和激發(squeeze-and-excitation networks,SE)模塊來提升分割精度。Guan等[8]在每個編碼器模塊中加入了SE網絡,并在每個解碼器中加入注意力引導過濾器(attention guide filter,AG)模塊,利用注意力機制對邊緣信息進行提取,設計出了3D注意力引導過濾器和壓縮-激發VNet(3D attention guide filter and squeeze-and-excitation VNet,3D AGSE-VNet)模型進行腦腫瘤圖像分割。

CNN在對數據特征進行提取時有出色的表現,但由于受感受野大小的限制,無法捕獲遠距離特征。因此,Vaswani等[9]提出視覺自注意力(vision transformer,ViT)模塊,作為基于注意力機制的模塊被引入到計算機視覺的各個領域,該模塊依靠其出色的遠距離特征捕獲和建模能力彌補了CNN的不足。隨著ViT模塊的引入,CNN-Transformer相互結合的模型被迅速提出。在腦腫瘤圖像分割方面,Wang等[10]提出使用Transformer進行多模態腦腫瘤圖像分割(brain tumor image segmentation using Transformer,TransBTS)模型,將注意力機制與CNN結合起來,構成編碼器-解碼器架構,利用3D CNN來捕獲空間和深度信息,通過ViT模塊對遠距離特征進行建模,最終得到高分辨率的分割結果。此外,Jia等[11]提出雙重Transformer的UNet(bi-Transformer U-Net,BiTr-UNet)模型,通過在編碼器-解碼器架構中的跳躍連接層上引入多個ViT模塊來提取更多數據特征信息,引入的多個ViT模塊在一定程度上提升了模型的分割性能,但模型自身參數量將會顯著提升,模型復雜度變高。Menze等[12]將20種最先進的腫瘤圖像分割算法應用于低級別和高級別神經膠質瘤,發現融合幾個算法產生的分割性能始終高于單獨的算法,說明腦腫瘤圖像分割算法性能還能進一步優化。

上述研究都在一定程度上提高了腦腫瘤圖像的分割精度,但這些模型在數據傳輸過程中由于池化操作會導致一些特征信息丟失,這將使分割模型在訓練數據中提取到的特征信息不足,最終應用在測試數據上效果不佳。此外,上述模型都未考慮模型的抗干擾能力,魯棒性不足。對此,提出融合雙重對抗學習的CNN-Transformer(CNN-Transformer model fusing dual adversarial learning,TransFDA)分割模型對腦腫瘤圖像進行分割,該模型在分割過程中使用輕量級CNN-Transformer模型,并通過引入雙重對抗學習來優化分割模型,以期在不增加模型參數量和復雜度的基礎上提高模型的分割精度和魯棒性。

1 模型的構建

TransFDA基于TransBTS和BiTr-UNet模型,整體訓練流程如圖1所示。由圖1可知,TransFDA的整體架構由3部分子模塊構成,分別為基礎分割模塊、判別器模塊、虛擬對抗訓練模塊。其中,基礎分割模塊由CNN-Transformer結構組成,在訓練過程中,首先將分割模型預測出的結果送入判別器模塊中,圖1中的黑色線條為基礎分割模塊正常生成預測標簽過程,紅色線條為預測標簽與訓練數據進行對抗過程;紅色實線代表2組數據分別進入判別器模塊進行特征匹配,生成判別差異損失,并將差異信息反饋給分割模型的過程;紅色虛線代表分割模型經判別模塊多次對抗優化后預測出越來越完整、準確的腫瘤區域;此過程直至判別模塊中生成的損失變化不大、收斂為止,此時分割模型達到了分割最優結果。另外,圖1中綠色線條為虛擬對抗訓練過程,其中實線部分為在原始圖像上增添干擾,然后將生成帶干擾的新圖像送入分割模型中,預測出帶干擾的標簽結果過程;虛線為將正常圖像預測出的分割結果和干擾圖像預測出的結果送入虛擬對抗訓練模塊,經過判別生成兩者之間的虛擬對抗損失的過程。最終,將上述2種損失與分割模型生成的損失結合起來,反饋到分割模型中。上述操作使得基礎分割模型在后續訓練過程中能夠處理之前分割過程中存在的不足,從而學習到圖像上更多的數據特征。

1.1 分割模型

基礎CNN-Transformer分割模型結構如圖2所示,其中淺藍色矩形代表特征圖,矩形周圍的數字代表每個特征圖的通道個數,紫色矩形代表在跳躍連接層上添加的ViT模塊,不同顏色箭頭代表不同的操作。由圖2可知,首先將4個模態的腦腫瘤圖像送入分割模型中,經一系列的特征學習,最終預測分割出圖像中病變的3塊區域。其中,編碼器模塊中一系列卷積和下采樣操作對進入模型中的高維數據進行降維處理,并在編碼器最后一層上使用空洞卷積來捕捉圖像的多尺度上下文信息,促使分割模型更好學習腦腫瘤圖像中的特征信息。為加強分割模型對圖像中重要特征信息的提取,每層卷積塊后都引入卷積注意力機制模塊(convolutional block attention module,CBAM),通過該模塊促使分割模型提取重要的空間和通道信息。另外,編碼器中每層提取到的特征都使用批歸一化(batch normalization,BN)操作,用整流線性單元(rectified linear unit,ReLU)激活函數進行處理,防止模型在訓練過程中出現梯度爆炸或梯度消失等問題,同時還能加快模型訓練。此外,解碼器同編碼器總體架構一致,都是由一系列卷積塊和上采樣操作組成;解碼器架構和編碼器架構相互對應,解碼器通過上采樣操作來恢復編碼器中下采樣操作壓縮的維度,融合層將提取到的特征與解碼器中對應的特征進行融合,以此來獲得更多的空間特征。最后,在解碼器最后一層使用卷積核大小為1×1的卷積和激活函數,將輸出的結果映射到期望分割的類別數。分割模型在最后2層跳躍連接處分別引入了1個ViT模塊來提取遠距離信息。

圖2 分割模型結構Fig.2 Segmentation model structure

1.2 判別器模塊

判別器模塊受馬爾可夫判別器啟發提出并使用。此模塊由3層3×3×3卷積塊組成,卷積塊之后添加BN與ReLU激活函數層,加快訓練速度。常用的判別器模塊在得到預測標簽后,通常使用預測標簽與真實標簽進行判別,而此判別器模塊受無監督與半監督模型啟發,不使用真實標簽,而是將分割模型預測到的結果與原始圖像送入到判別器模塊中,在判別器處理過程中,整個腦部都被認作存在腫瘤的區域。首先,將整個腦部與預測得到的標簽一同送入判別器模塊中進行判別,由于剛開始分割模型能力有限,預測到的分割標簽上腫瘤區域較小,判別器很容易就對其判斷為假,并生成損失,同時將信息反饋給基礎分割模型,此時基礎分割模型對整個腦部中之前未分割的區域加強識別分割;隨著多次訓練的疊加,基礎分割模型與判別器模塊反復相互優化,直至判別模塊中的損失達到擬合,不再發生過多變化為止,此時模型會分割出最優結果。

1.3 虛擬對抗訓練

在訓練過程中訓練數據數量有限,分割模型魯棒性存在不足,對外界可能出現的各種干擾應對能力較差。為提高模型抵抗干擾數據的能力,引入虛擬對抗訓練。首先,找到使得輸出偏差最大的擾動方向,在此方向上對輸入的原始數據加入擾動,該擾動通過分割模型的損失與正則化后的原始數據梯度組合而成。然后,將原始數據與添加擾動后的新數據送入分割模型,經過處理后分別得到各自的分割結果標簽,其中添加擾動的數據分割出的標簽不同于正常分割結果,會產生異常標簽。同上述判別器模塊一致,此模塊中同樣不使用真實標簽,將添加擾動后分割得到的標簽與正常數據分割后得到的標簽進行判別,將得到這2個標簽之間的損失,并將其反饋給分割模型;分割模型收到多次信息反饋優化后,添加干擾的數據分割出的結果與正常數據經分割模型預測出的標簽越來越一致。在整個對抗過程中,隨著訓練輪次的疊加,模型魯棒性逐漸加強,外界干擾對分割模型的影響越來越低,模型將變得更加穩健。

2 實驗與結果

2.1 實驗數據集

分割模型訓練和驗證過程中使用國際醫學圖像計算和計算機輔助干預協會(medical image computing and computer assisted intervention society,MICCAI)提供的Brats2020多模態腦腫瘤數據集[12]。此數據集將369張病例作為訓練集,125張病例作為驗證集,訓練集中包含高級別膠質瘤(high-grade glioma,HGG)和低級別膠質瘤(low-grade glioma,LGG)信息,驗證集中案例級別未知。由于Brats2020公開數據集中未提供測試集數據,為進一步驗證模型性能,從Brats2021數據集中隨機收集125例未出現在Brats2020數據集中的案例作為測試集。每個案例的4個模態和真實標簽如圖3所示。由圖3可知,4個模態分別為液體衰減反轉恢復序列(fluid attenuated inversion recovety,FLAIR)、T1、T1ce、T2,訓練數據集中提供了真實分割標簽,包括背景、非增強腫瘤區域(non-enhancing tumor,NET)、浮腫區域(edema,ED)和增強腫瘤區域(enhancing tumor,ET)。標簽圖像中的綠色部分代表ED,此區域對應標簽2;紅色代表NET,此區域對應標簽1;黃色代表ET,此區域對應標簽4;背景區域對應標簽0。低級別膠質瘤相較高級別膠質瘤會缺少一些病變區域,若模型分割到這些區域時分割系數會顯示為0,因此對低級別膠質瘤進行分割時難度會有所上升。最終任務是對腦腫瘤圖像中的腫瘤整體區域(whole tumor,WT)、腫瘤核心區域(tumor core,TC)和增強腫瘤區域(enhancing tumor,ET)進行分割。其中WT區域包含所有的標簽,即標簽圖像中紅色、綠色、黃色3個區域的合并部分,TC區域為紅色、黃色區域的合并部分,ET為標簽中黃色區域。

(a) FLAIR圖像 (b) T1圖像 (c) T1ce圖像 (d) T2圖像 (e) 真實標簽圖3 不同模態的腦腫瘤和標簽Fig.3 Different modality brain tumor images and labels

2.2 數據集預處理

數據集中所有病例包含的4個模態數據大小都為240×240×155,4個模態共享分割標簽。首先,將4張不同模態腦腫瘤圖像進行合并處理,將其合并成1張4D(H×W×D×C,C=4)圖像,并與分割標簽一同保存為pkl文件。這是由于此格式的文件讀寫速度較快,有利于加快訓練速度。在對實驗數據進行觀測時發現,原始圖像的周圍存在著大規模的黑邊,因此將實驗數據統一裁剪成128×128×128尺寸大小的圖片,并對其在[-0.1,0.1]范圍內隨機變換強度,在[0.9,1.1]的范圍內縮放。

2.3 實驗環境與配置

實驗使用Ubuntu16.04操作系統,在2張帶有24GB顯存的NVIDIA GeForce RTX 3090顯卡上進行模型的訓練、驗證、測試。訓練過程中使用初始化學習率為0.0002的Adam優化器。實驗基于Pytorch框架實現,批處理大小為8,共迭代1000次。

2.4 評價指標

為了對所提模型進行性能評估,采用95%豪斯多夫距離(Hausdorff distance,HD)和戴斯相似系數(Dice similarity coefficient,DSC)作為分割模型的評價指標。

HD指標在醫學圖像分割中是指2張掩碼邊界像素距離集合中的最大值,即真實分割標簽與預測出的標簽2組點集之間距離的度量,具體計算公式如下所示:

(1)

其中,y′和y分別表示模型預測標簽和真實標簽,a和b為2個體素集中的體素點,H為HD指標的計算結果,值越小,說明真實分割標簽與預測標簽兩組點集之間距離越近,相似度就越高。

DSC指標為語義分割中的重要評價指標之一,可以衡量2個樣本之間的相似程度;數值越接近1,表明模型分割精度越好,具體計算公式如下所示:

(2)

其中,PDSC為DSC指標的計算結果,Dtrue代表分割圖像的真實標簽,Dpre代表模型分割出的預測標簽。

2.5 損失函數

分割模型中使用廣義戴斯損失(generalized Dice loss,GDL)函數。GDL函數是在戴斯損失函數基礎上優化而來,后者是針對特定類別分割的損失,在需要分割出多個病灶的場景下并不適用,而GDL函數可將多場景下分割的病灶損失整合為1個指標,將其作為分割結果進行評估,具體計算方式如下所示:

(3)

其中,PGDL為GDL函數計算結果,Rln表示第l個類別圖像在第n個位置的真實分割結果,Pln表示相應的預測概率,Wl表示第l個類別的權重。

在判別器模塊中,為了使最終分割結果更加接近于真實標簽結果,使用和大多數生成式對抗模型中相同的策略,最大化判別器判別能力,最小化分割模型中分割結果與真實圖像之間的差異性,具體可表示為:

(4)

其中,θs代表基礎分割模塊的參數,θd代表判別模塊的參數,X為輸入的數據集合,包括輸入圖像和生成的預測標簽。

判別器模塊對預測出的標簽與整個腦部之間的差異大小進行判別。判別差異損失定義為

(5)

其中,Ladv(θs;θd;X)為判別器模塊生成的判別差異損失,x為輸入圖像,E(x,y′)為輸入圖像和生成預測標簽的集合,H、W為圖片高度與寬度,c、d分別為圖片高度與寬度上的像素點,ψ∈(0,1)H×W的大小表示判別器的功能;當從分割模型中生成分割掩碼時,η=0,當從輸入圖片中提取樣本時,η=1。

虛擬對抗損失用來判別預測出的分割標簽與添加擾動后生成的標簽之間的差異,最小化虛擬對抗損失,即是最大化模型克服添加擾動數據的能力,使模型的魯棒性得以加強。具體定義如下所示:

(6)

其中,Lvat(θs;X;radv)為虛擬對抗訓練生成的虛擬對抗損失,radv是對原始圖像上施加的擾動信息,DKL為KL(Kullback-Leibler)散度,F代表分割模型。

整個模型最終損失由上述3部分損失函數線性累加組成,具體定義如下所示:

L(θ;X;Y)=PGDL(θs;Y)+λLadv(θs;θd;X)+ωLvat(θs;X;radv),

(7)

其中,L(θ;X;Y)為最終模型的總損失,θ代表模型中所有模塊參數,Y為真實標簽結果與預測標簽結果集合,λ、ω分別為2組對抗模塊對應的超參數。為保證最終分割模型的穩定性,經過多次的參數測試后,最終設置為λ=0.002、ω=0.003,其他參數值的設置可能會出現異常結果。模型訓練過程中損失函數的數值變化情況如圖4所示。由圖4可知,隨著訓練輪次的增加,損失函數值逐漸下降,在訓練到900輪左右時,損失函數的變化趨勢已經不大,呈現出收斂狀態。

圖4 損失函數變化趨勢Fig.4 Change trend of loss function

2.6 TransFDA實驗結果

TransFDA在Brats2020驗證集上的分割結果如表1所示。由表1可知,對ET、WT、TC的PDSC平均值分別達到0.7530、0.8922、0.7909,對ET、WT、TC的H平均值分別達到27.51、6.32、15.67;其中,WT的分割結果較ET與TC更優,這是由于WT在腦腫瘤中所占比例較大,ET與TC病變區域較小,另外這2個區域數據標簽之間存在較大的不確定性,尤其是對低級別腦腫瘤數據進行分割時,模型獲取信息不足,導致最終結果較差。

表1 Brats2020驗證集上的分割結果Tab.1 Segmentation results on the Brats2020 validation dataset

2.7 對比實驗

為驗證TransFDA的性能優勢,將其分別與不同的2D、3D腦腫瘤分割模型和參加BraTS2020挑戰賽的其他團隊所提模型進行對比,結果如表2、表3所示,表格中結果來源于各種模型原論文和挑戰賽官方網站上提供的實驗數據。由表2、表3可知,所提模型對ET、WT、TC的PDSC數值均高于其他腦腫瘤分割模型,所提模型對ET、WT、TC的H也取得了較為優秀的結果。

表2 對比實驗結果Tab.2 Comparison of experimental results

表3 與Brats2020挑戰賽其他團隊在驗證集上的對比實驗結果Tab.3 Comparison of experimental results with other teams in the Brats2020 Challenge on the validation dataset

此外,TransFDA與經典CNN-Transformer分割模型參數的對比結果如表4所示。由表4可知,相較于經典的CNN-Transformer分割模型,TransFDA在模型參數量、計算量、模型大小上有著明顯的優勢。

表4 與經典CNN-Transformer分割模型的參數對比Tab.4 Comparison of parameters with classic CNN-Transformer segmentation models

2.8 消融實驗

在Brats2020驗證集和Brats2021測試集上進行了5組消融實驗,以此來驗證提出的模塊是否能夠優化模型的分割性能,所有消融實驗均使用相同的實驗環境和超參數,驗證集上的實驗結果如表5所示。由表5可知,通過對比前3個模型實驗數據發現,在驗證集上基礎分割模型相較于TransBTS和BiTr-UNet對WT、TC的分割精度有明顯提升,特別是對TC,相較于TransBTS,PDSC數值提升了0.0230,H數值降低了6.68。對比第3個模型和第4個模型實驗數據可以看出在添加新型判別器模塊后,模型對ET、WT分割精度進一步得到了提升,但對TC分割精度有所下降,通過對Brats2020實驗數據的分析,最終判斷是因為分割模型對TC的過分割以及模型缺乏抗干擾能力所致。因此,再引入虛擬對抗模塊(即完整的TransFDA)來進一步加強模型分割能力。對比第4個模型和第5個模型實驗數據可以看出添加虛擬對抗訓練模塊后,模型對3個區域分割結果都得到了有效的提升,其中對ET、WT的分割結果達到了最優,相較于基礎模型,PDSC數值分別提升0.0141、0.0123,H數值分別降低了3.29、1.88。另外,對TC的分割性能相較于基礎分割模型有所不足,通過對所有數據參數進行分析,推測是因為雙重對抗模塊對Brats2020驗證集中的一些數據存在過分割;但綜合所有參數指標的比較結果,最終表明TransFDA性能仍能達到最優。

表5 在驗證集上的消融實驗Tab.5 Ablation experiment on validation set

測試集上的實驗結果如表6所示。由表6可知,在基礎分割模型上引入雙重對抗學習后,TransFDA對ET的H數值有輕微提升,表示圖像真實標簽和模型預測生成標簽的相似度降低,其余所有指標相較于基礎分割模型都有著明顯的性能提升;對于ET、WT、TC的PDSC平均值分別達到0.8438、0.9155、0.8900,對ET、WT、TC的H平均值分別達到2.00、5.00、2.83;相較于表格中其他CNN-Transformer分割模型,PDSC數值都有著一定的提升,但比較H數值,發現模型的性能有一定的欠缺。

表6 在測試集上消融實驗Tab.6 Ablation experiment on the test set

2.9 定性與定量實驗

為更加直觀地分析模型中各模塊的分割效果,分別挑選了高級別膠質瘤和低級別膠質瘤數據進行可視化處理,通過可視化結果能夠更好地觀測各模塊之間的差異。MRI腦腫瘤水平面、矢狀面、冠狀面分割結果如圖5、圖6所示。由圖5、圖6可知,在分割區域上可以明顯看出完整的TransFDA分割出的區域較其他模型更加豐富。

圖5 MRI高級別腦腫瘤(validation_058)分割結果在水平面、矢狀面、冠狀面的對比Fig.5 MRI high-grade brain tumor (validation_058) segmentation results in the horizontal,sagittal,and coronal planes compared

由于Brats2020驗證集數據不提供真實標簽,因此將各模型的指標結果繪成表格,結果如表7、表8所示。由表7可知,對于高級別膠質瘤的分割,TransFDA對ET、WT、TC的分割性能提升顯著,而高級別膠質瘤較易分割;對于難分割的低級別膠質瘤,由表8可知,DSC和HD指標也有一定提升,且不同模型性能差異較為明顯,表示所提TransFDA對于此類型數據的分割效果較好。

表7 不同模塊在Brats2020_validation_058上的實驗結果Tab.7 Experimental results of different modules on Brats2020_validation_058

表8 不同模塊在Brats2020_validation_072上的實驗結果Tab.8 Experimental results of different modules on Brats2020_validation_072

同時,在Brats2021測試集上進行可視化實驗,測試集數據提供真實標簽,可視化結果如圖7所示。由圖7可知,TransFDA分割出來的標簽和真實標簽相似度極高,在分割區域的細節處理上更加完善,對腦腫瘤邊緣區域的分割相較于其他模型更加細致。以上實驗結果表明,對于較易分割的高級別腦腫瘤數據,大多數模型分割性能差異表現不大,但對于較難分割的低級別腦腫瘤數據,不同模型的性能差異較為明顯。通過以上定性與定量結果分析可知,TransFDA模型的綜合性能表現最優。

圖7 MRI高級別腦腫瘤測試集分割結果在水平面、矢狀面、冠狀面的對比Fig.7 MRI high-grade brain tumor test set segmentation results in the horizontal,sagittal,and coronal planes compared

3 結論

提出融合雙重對抗學習的CNN-Transformer分割模型TransFDA,所提TransFDA是對傳統CNN-Transformer模型進一步優化改進而來,加入的新型判別器模塊能有效加強基礎分割模塊對于腦腫瘤圖像上特征信息的提取,顯著提升模型分割性能;另外,針對基礎分割模型中存在的過分割和抗干擾能力不足問題,引入虛擬對抗訓練模塊,可增強分割模型的魯棒性;最后,在Brats2020驗證集上進行性能測試,對ET、WT、TC的PDSC平均值分別達到0.7530、0.8922、0.7909,對ET、WT、TC的H平均值分別達到27.51、6.32、15.67。實驗結果表明,TransFDA可以有效提升分割精度,同時與經典的基于CNN和CNN-Transformer的2D、3D分割模型,以及Brats2020挑戰賽上其他團隊所提模型相比較,在腦腫瘤分割精度上都有著一定的優勢。通過引入Brats2021上部分數據集來進一步證明在測試集上TransFDA的性能優勢,發現對ET、WT、TC的PDSC平均值分別達到0.8438、0.9155、0.8900,對ET、WT、TC的H平均值分別達到2.00、5.00、2.83。雖然所提TransFDA在腦腫瘤分割上有優秀表現,但仍然存在一些問題,比如對于低級別膠質瘤的分割,在增強腫瘤與腫瘤核心區域的分割精度較為不足,特別對腫瘤邊界的分割較為毛糙,不夠平滑,并存在一些過分割問題,后續將針對這些問題展開深入研究。