戴昂 宋亞男 徐榮華 方俞澤
摘要:文章針對形狀復(fù)雜、邊界模糊的腦腫瘤難以實現(xiàn)精確分割的問題,提出一種基于卷積注意力機制和Transformer多頭注意力機制的U型分割網(wǎng)絡(luò)。文章首先設(shè)計了基于通道注意力和空間注意力的卷積模塊,提高了模塊對局部關(guān)鍵特征的提取能力:其次使用一種結(jié)構(gòu)更精簡的Transformer模塊作為網(wǎng)絡(luò)的瓶頸層,利用其多頭注意力機制對全局特征進行充分感知:最后在BraTS 2021數(shù)據(jù)集上進行了實驗。實驗結(jié)果表明文章算法在增強腫瘤區(qū)域、腫瘤核心區(qū)域和整個腫瘤區(qū)域的Dice系數(shù)評分分別為87. 51%,90. 69%和93. 47%,可以有效提高腦腫瘤分割精度。
關(guān)鍵詞:腦腫瘤分割:混合網(wǎng)絡(luò):Transformer:注意力機制
中圖分類號:TP391
文獻標(biāo)志碼:A
0 引言
在臨床中,通過核磁共振影像( Magnetic ResonanceImaging,MRI)對腦腫瘤精確分割是醫(yī)生進行腦腫瘤診斷和制定治療方案的關(guān)鍵環(huán)節(jié),這個過程主要需要大量的專業(yè)經(jīng)驗,容易出錯而且費時費力,而傳統(tǒng)的基于閾值的分割方法應(yīng)用到腦腫瘤分割時有較大的局限性。
近年來,卷積神經(jīng)網(wǎng)絡(luò)( Convolutional Neural Net-works.CNN)因其強大的特征提取能力而被廣泛應(yīng)用于圖像處理任務(wù)。在醫(yī)學(xué)圖像分割領(lǐng)域,Ronneberger等[1]首次提出了U-Net架構(gòu),實現(xiàn)了端到端的自動分割,該架構(gòu)迅速成了醫(yī)學(xué)圖像分割最有效的方法。為了充分利用MRI等三維數(shù)據(jù)的多個維度的信息.CiCek等[2]將U-Net擴展成三維結(jié)構(gòu),提出了3 D U-Nel,該網(wǎng)絡(luò)在腦腫瘤數(shù)據(jù)集上分割的各項分割指標(biāo)均優(yōu)于二維的U-Net網(wǎng)絡(luò)。隨著注意力機制的有效性被證實,研究人員逐漸將其引入以提高醫(yī)學(xué)圖像的分割精度。例如,郝曉宇等[3]將位置注意力和通道注意力引入3 D U-Net.提出DAU-Net用于肺腫瘤分割,有效地提升了肺腫瘤的分割精度。隨著研究的深入,在自然語言處理領(lǐng)域取得了成功的一種深度神經(jīng)網(wǎng)絡(luò)Transformer逐漸被應(yīng)用于圖像的分類、檢測和分割任務(wù)。Hacamizadeh等[4]提出針對三維醫(yī)學(xué)圖像分割的網(wǎng)絡(luò)UNelFormer.該網(wǎng)絡(luò)依然基于U-Nec框架,但編碼器和解碼器都利用Transformer實現(xiàn),僅依靠Transformer強大的長距離依賴表示能力,實現(xiàn)了對腦腫瘤的精確分割。
雖然在模型中引入卷積注意力或Transformer的多頭白注意力的方法在三維醫(yī)學(xué)圖像分割中取得了一定的效果,但如何結(jié)合CNN的局部特征提取能力和Transformer的全局特征提取能力,更有效地發(fā)揮注意力機制的作用,提升三維醫(yī)學(xué)圖像的分割精度,仍然值得研究。
綜上,本文提出了一個融合CNN通道注意力、空間注意力和Transformer多頭注意力的腦腫瘤分割算法。實驗結(jié)果表明,本文算法在卷積編碼塊中合理地引入了注意力機制,并有效融合了Transformer架構(gòu),能夠充分發(fā)揮CNN和Transformer的優(yōu)點,提升腦腫瘤分割精度。
1 本文方法
1.1 模型整體結(jié)構(gòu)
本文提出的腦腫瘤分割模型整體結(jié)構(gòu)如圖l所示,模型輸入尺寸為HxWxDx4.其中H.W和D表示三維圖像的3個維度,4表示輸入數(shù)據(jù)的4種模態(tài)。
模型由編碼器和解碼器構(gòu)成。編碼器部分首先堆疊了3個具有相同結(jié)構(gòu)的卷積層,這些卷積層都由卷積下采樣、通道注意力編碼模塊和空間注意力編碼模塊組成。編碼器的底層瓶頸層則由卷積下采樣和Transformer模塊構(gòu)成。解碼器將編碼器每一層的下采樣卷積用上采樣轉(zhuǎn)置卷積替換,用于恢復(fù)特征圖的分辨率,然后對來自編碼器部分的同分辨率特征圖進行通道維度的拼接,其余部分和編碼器部分保持一致。最后的分割層的作用是將特征圖恢復(fù)為與輸入同等大小,然后由卷積核大小為lxlxl的卷積將通道數(shù)量恢復(fù)為4.經(jīng)過Softmax激活函數(shù)得到分類輸出。
1.2 通道注意力層
編碼器和解碼器的卷積編碼模塊都由通道注意力層和空間注意力層組成。其中,通道注意力層首先使用兩個卷積核為1x1x1的點卷積和一個卷積核為3x3x3的深度可分離卷積進行特征提取,與使用兩個卷積核為3x3x3的普通卷積相比,計算量顯著降低,并且歸一化方式選擇Croup Norm,激活函數(shù)選擇Leaky Relu;然后在卷積特征提取部分后面引入了經(jīng)典通道注意力模塊SE[5].強化某些重要通道的特征圖權(quán)重:最后在SE后面拼接了一個由卷積實現(xiàn)的全連接模塊,加強全局特征感知能力。通道注意力層的結(jié)構(gòu)如圖2所示。
1.3 空間注意力層
空間注意力層用CBAM[6]中的空間注意力模塊代替了通道注意力層中的通道注意力模塊,使其更加關(guān)注于特征圖的空間位置關(guān)系,而不是對通道特征的強化。對于輸入尺寸為H×W×DxC的特征圖,空間注意力模塊首先對其進行通道維度上的最大值池化( MaxPool)與平均值池化(AvgPool),得到兩個H×WxDxl的特征圖,然后將兩個特征圖按通道方向拼接,接著經(jīng)過一個7x7x7的卷積并使用Sigmoid函數(shù)進行激活,最后將激活函數(shù)輸出的特征權(quán)重矩陣與輸入特征圖相乘,得到對感興趣的區(qū)域進行了增強并抑制了無用信息后的特征圖。
1.4
Transformer模塊
與卷積神經(jīng)網(wǎng)絡(luò)相比.Transformer具有動態(tài)注意、全局上下文感知和更好的泛化能力等優(yōu)點。為了充分利用其多頭注意力機制的全局特征編碼能力,并減少參數(shù)量,受Wu等[7]提出的卷積視覺Transformer方法的啟發(fā),本文使用卷積對原始的Transformer進行簡化。簡化的Transformer模塊結(jié)構(gòu)如圖4所示。
簡化后的Transformer模塊沒有線性位置編碼投影層,而是直接通過一個深度可分離卷積.生成多頭注意力所需的p,K,V矩陣。具體的,對于經(jīng)過4次
下采樣的特征圖
,通過卷積核大小為3x3x3的深度可分離卷積實現(xiàn)卷積投影,然后將特征圖在通道維度展平,并使用Layer Norm進行歸一化,得到Q,K,V矩陣。接著進行多頭注意力機制的計算,算法流程可以描述為:
其中,z1表示第l∈[1,2.¨.,L]個Transformer層的輸入,z0為初始輸入E…,LN表示Layer Norm.MLP表示全連接模塊,其結(jié)構(gòu)也與原始的Transformer中的不同,使用的是與通道注意力層和空間注意力層中的全連接模塊一樣的結(jié)構(gòu)。MHSA表示多頭自注意力,計算公式為:
2 實驗過程與結(jié)果分析
2.1 實驗數(shù)據(jù)與預(yù)處理
本文的實驗數(shù)據(jù)由腦腫瘤分割挑戰(zhàn)賽BraTS202l[s]提供,數(shù)據(jù)集由4種MRI模態(tài)組成,分別是Tlw.Tlcew,T2w和Flair。每種模態(tài)的數(shù)據(jù)大小為240x 240×155.腫瘤分割的區(qū)域分為增強腫瘤區(qū)域ET、腫瘤核心區(qū)域TC和全部的腫瘤區(qū)域WT。將1 251個數(shù)據(jù),按照4:1的比例劃分為訓(xùn)練集和測試集。在訓(xùn)練過程中,本研究將數(shù)據(jù)隨機裁剪成128×128x128的大小,并進行了歸一化處理。
2.2模型訓(xùn)練
實驗中的模型均在Uhuntu 20. 04.1操作系統(tǒng)下,基于PyTorch l.10.0實現(xiàn),使用NVIDIA 3090 CPU,以batchsize為2,訓(xùn)練500輪,訓(xùn)練的優(yōu)化器選擇Adam.學(xué)習(xí)率設(shè)置為0.000 5,訓(xùn)練的損失函數(shù)選擇Dice損失。
2.3評價指標(biāo)
為了定量比較和評估分割結(jié)果,本文計算了醫(yī)學(xué)圖像分割領(lǐng)域中廣泛使用的評價指標(biāo)Dice系數(shù),計算公式如下:
其中,|Xn Y|表示樣本X和樣本l,的交集.|X|和| Y|為樣本元素數(shù)量。
2.4 實驗結(jié)果與分析
將本文算法在測試集上與CNN方法3 D U-Net、Transformer方法UNetFormer和CNN與Transformer的混合方法TransBTS[9]進行了分割性能的比較。如表1所示,實驗結(jié)果表明,本文算法在ET.TC和WT的平均Dice系數(shù)與其他算法相比都取得了最高的分數(shù)。
如圖5所示為各方法的分割結(jié)果對比圖片,圖中從左到有依次為3 D UNet、UNetFormer、TransBTS、本文算法和真值。
從圖中可以觀察到,本文提出的模型在形狀不規(guī)則且模糊重疊的腦腫瘤上,分割更準(zhǔn)確,更完整。例如在第一個腦腫瘤數(shù)據(jù)上,與本文算法相比,其他方法都存在更加嚴重的過度分割,在第二個腦腫瘤數(shù)據(jù)上.3 D UNec存在較大的分割錯誤,從整體與細節(jié)對比中可以發(fā)現(xiàn),本文算法與真實值最為接近。
2.5 消融實驗
為了驗證本文算法中卷積特征提取模塊選擇的通道注意力和空間注意力的有效性以及Transformer模塊的有效性,進行了各模塊的消融實驗,消融實驗中作為對比的模型也以相同的數(shù)據(jù)和參數(shù)訓(xùn)練得到。實驗結(jié)果如表2所示。
從實驗結(jié)果可以看出:
(l)不使用Transformer模塊,僅依靠本文設(shè)計的通道注意力層和空間注意力層構(gòu)建基礎(chǔ)分割網(wǎng)絡(luò),也能取得比3 D U-Net分割效果更好的分割結(jié)果,表明本文設(shè)計的基礎(chǔ)框架是有效的。
(2)直接用CBAM雙注意力模塊代替通道和空間注意力層中的注意力機制時.Dice評分與最終模型相比有所下降,可能是使用CBAM模塊直接計算完通道和空間注意力,會使得網(wǎng)絡(luò)在訓(xùn)練時過于關(guān)注某些特征,從而更加容易過擬合。
(3)將SE模塊替換成CBAM中的通道注意力模塊后,分割性能有所提升,但依然低于本文最終的方法,表明將通道注意力模塊和空間注意力模塊嵌入不同的卷積編碼層分別強化通道和空間特征,能夠提升網(wǎng)絡(luò)性能,實現(xiàn)分割精度的提升,而且本文使用的注意力機制組合方式更加合理,能更加有效地發(fā)揮CNN和Transformer各自的優(yōu)點。
3 結(jié)語
本文提出了一個用于三維腦腫瘤圖像分割的網(wǎng)絡(luò),該網(wǎng)絡(luò)基于通道注意力、空間注意力、白注意力構(gòu)建。網(wǎng)絡(luò)的淺層使用卷積注意力模塊有助于對圖像的局部空間信息進行充分編碼。瓶頸層使用簡化的Transformer模塊進行全局特征提取,在引入少量參數(shù)的情況下,充分發(fā)揮Transformer結(jié)構(gòu)優(yōu)秀的動態(tài)注意能力和學(xué)習(xí)全局語義相關(guān)性的能力。實驗結(jié)果表明,在腦腫瘤數(shù)據(jù)集上,本文提出的算法取得了更高的Dice系數(shù)評分,可以更準(zhǔn)確地分割腫瘤并有助于臨床診斷的改善。考慮到模型復(fù)雜度也是算法能否實際應(yīng)用的重要指標(biāo),后續(xù)將在保證分割精度的情況下,在模型輕量化策略上做出進一步改進。
參考文獻
[1]RONNEBERGER 0, FISCHER P, BROX T. U-net: Convolutional networks for biomedical imagesegmentation : international conference on medical imagecomputing and compurer-assisted inlervention [ C lBerlin: Springer, Cham, 2015.
[2] qICEK O, ABDULKADIR A, LIENKAMP S S. etal. 3D U-Net: leaming dense volumetric segmentationfrom sparse annotation: International conference onmedical imag
[3]郝曉宇 ,熊俊峰 ,薛旭東 ,等 .融合 >-;7\注意力機制3D U -Net的肺腫瘤分割 [ J] .中國圖象圖形學(xué)報,2020( 10) :2119-2127.
[4]HATAMIZADEH A. XU Z. YANCJ D, et al.UNetFormer: A unified vision transformer model andpre-training framework for 3D medical imagesegmentation[ EB/OL ] . ( 2022 - 04 - 05 )
[ 2023 - 02 -14 ] .
https : //arxiv. org/ab s/2204. 0063 1.
[5] HU J, SHEN L. SUN G. Squeeze -and -excitaLionnetworks: Proceedings of the IEEE conference oncomputer vision and pattern recognition [ C ] . Sah LakeCity: IEEE, 2018.
[6]WOO S, PARK J, LEE J Y, et al. Cbam:Convolutional block attention module : Proceedings of theEuropean c.onference on computer vision ( ECCV) [ C] .Berlin: Springer, Press. 2018.
[7]WU H, XIAO B, CODELLA N, et al. Cvt:Introducing convolutions to vision transformers:Proceeding{s of the IEEE/CVF international conferenceon compurer vision[ C] . Salt Lake City : IEEE , 2021.
[8]BAID U, C.HODASARA S, MOHAN S, et al. Thersna-asnr-mic-.cai brats 2021 benc-.hmark on brain tumorsegmentation and radiogenomic classification [ EB/OL] .( 2021- 09 - 12 )
[ 2023 - 02 - 14 ] . https ://arxiv. org/abs/2107. 02314.
[9]WANG W, CHEN C, DING M. et al. Transbts:Multimodal brain tumor segmentation using transformer:International conferenc-.e on medical image computing andcomputer-assisted intervention [ C ] . Berlin : Springer,Cham , 2021.
(編輯王雪芬 )