基于Transformer的核磁共振腫瘤圖像分割研究

呂嫄

(蕪湖職業(yè)技術(shù)學(xué)院 基礎(chǔ)教學(xué)部,安徽 蕪湖,241000)

核磁共振成像可以很好地顯示人體的組織信息,用于區(qū)分腫瘤是否為癌性以及癌細(xì)胞是否轉(zhuǎn)移。因此,核磁共振成像可以輔助醫(yī)務(wù)人員規(guī)劃腫瘤治療方案。近年來,隨著深度學(xué)習(xí)技術(shù)在圖像處理領(lǐng)域的迅速發(fā)展,運(yùn)用深度學(xué)習(xí)技術(shù)對腫瘤病人的核磁共振影像進(jìn)行自動分割可以更加快速準(zhǔn)確地為醫(yī)生提供診斷依據(jù)。

基于深度學(xué)習(xí)技術(shù)的卷積神經(jīng)網(wǎng)絡(luò)被廣泛應(yīng)用于醫(yī)學(xué)圖像分割領(lǐng)域[1-3]。U-net是最早提出U形分割架構(gòu)的卷積神經(jīng)網(wǎng)絡(luò),其設(shè)計思想得到了廣泛的應(yīng)用[4-6]。在卷積神經(jīng)網(wǎng)絡(luò)中,有很多工作都證明了多尺度特征對于圖像分割任務(wù)非常重要[7-8]。池化操作會損害這種多尺度的信息,空洞卷積解決了這個問題[9]。但是這種卷積會造成信息的不連續(xù),因此,本文提出一種不使用卷積的模型。

Transformer最早是用于解決自然語言處理領(lǐng)域中無法并行訓(xùn)練的問題。直到視覺Transformer巧妙地將其應(yīng)用于圖像處理,彌補(bǔ)了卷積神經(jīng)網(wǎng)絡(luò)的一些缺點,推進(jìn)了醫(yī)學(xué)圖像處理領(lǐng)域的發(fā)展[10-11]。Swin-Transformer在視覺Transformer基礎(chǔ)上提出了滑動窗口算法,降低了計算復(fù)雜度[12]。nnFormer和VT-UNet是兩個完全基于Transformer架構(gòu)的3D分割網(wǎng)絡(luò)[13-14]。TransUNet是第一個提出結(jié)合卷積和Transformer的醫(yī)學(xué)圖像分割網(wǎng)絡(luò),但是它只用了一個卷積層來提取Transformer的特征[15]。ViTAE通過使用不同膨脹率的多重卷積向下采樣,將輸入圖像嵌入到具有多尺度上下文的令牌中[16]。UNETR,TransBTS都是基于視覺Transformer架構(gòu)的3D醫(yī)學(xué)圖像分割網(wǎng)絡(luò),計算復(fù)雜度高[17-18]。Cross-view Transformer通過在全局池化層之前鏈接特征圖,在視圖之間傳輸信息,在胸部X射線數(shù)據(jù)集上取得優(yōu)越的性能[19]。FCT(fully convolutional Transformer)是第一個用于醫(yī)學(xué)影像應(yīng)用的全卷積Transformer模型。它利用Transformer提取圖像的長距離語義依賴,然后利用卷積提取圖像的局部屬性,在ISIC 2017數(shù)據(jù)集上取得較好性能[20]。D-former 提出了一種膨脹Transformer,在不增加計算成本的情況下擴(kuò)大感受野,提高了模型對于3D醫(yī)學(xué)圖像的分割性能。DS-TransUNet 將分層Swin Transformer 融合到U形結(jié)構(gòu)的編碼器與解碼器中,有效地提取了非局部依賴和多尺度上下文信息,提高了不同醫(yī)學(xué)圖像的語義分割質(zhì)量。GFENet以視覺Transformer作為基礎(chǔ)單元構(gòu)建了一個金字塔結(jié)構(gòu)的特征提取器,通過線性預(yù)測和局域解碼對提取的特征進(jìn)行局部增強(qiáng),并采用全局注意力將位置信息嵌入特征,在CVC-ColonDB數(shù)據(jù)集上獲得了優(yōu)于其他高性能分割模型的結(jié)果。經(jīng)過綜合分析各模型的優(yōu)缺點,本文提出了一種完全基于Transformer的分割網(wǎng)絡(luò),它以Transformer的變體Swin Transformer作為基礎(chǔ)模塊,不僅可以準(zhǔn)確地分割三維核磁共振的腦腫瘤圖像,而且參數(shù)量少,計算復(fù)雜度低。

1 方法

1.1 數(shù)據(jù)集

本文為了驗證所提出網(wǎng)絡(luò)結(jié)構(gòu)的有效性,使用了公共的腦腫瘤分割數(shù)據(jù)集BraTS2021進(jìn)行了實驗驗證[21]。該數(shù)據(jù)集一共包含1 251例患者的三維核磁共振影像。本文對所有樣本進(jìn)行隨機(jī)抽樣,將834例樣本作為訓(xùn)練集,208例樣本作為驗證集,209例樣本作為測試集。其中,每例樣本由4種模態(tài)(Flair,T1,T1CE,T2)的3D圖像以及1個共享標(biāo)簽組成。

每個模態(tài)圖像都是磁共振成像,長×寬×高為240×240×155,單個的體積像素是1 mm3。如圖1所示,T1(T1-weighted imaging)通過對人體進(jìn)行磁共振成像,主要用于顯示解剖結(jié)構(gòu);T1CE(T1-weighted imaging with contrast enhancement)在T1序列的基礎(chǔ)上加入了造影劑,主要用于顯示腫瘤和炎癥等異常組織;T2(T2-weighted imaging)使用對T2信號敏感的磁共振成像方法,用于顯示病變的水腫或液體;Flair(Fluid Attenuated Inversion Recovery)通過將T2序列與T1序列結(jié)合,用于顯示腫瘤、炎癥和其他異常組織。

圖1 數(shù)據(jù)集4種模態(tài)示例Fig.1 Four modal examples for datasets

1.2 網(wǎng)絡(luò)結(jié)構(gòu)

本文提出的分割網(wǎng)絡(luò)的整體結(jié)構(gòu)如圖2所示,它主要由收縮模塊、擴(kuò)張模塊和跳轉(zhuǎn)連接組成。其中收縮模塊用于提取和融合輸入圖像的特征;擴(kuò)張模塊用于恢復(fù)特征圖的原始分辨率;跳轉(zhuǎn)連接有助于恢復(fù)圖像的細(xì)粒度細(xì)節(jié)。當(dāng)輸入一張圖像后,經(jīng)過多個編碼層得到不同尺度的立體特征圖(圖2中輸入分支中的立方體),并將編碼后的特征圖直接輸入擴(kuò)張模塊進(jìn)行解碼得到新的特征圖(圖2中輸出分支中的立方體)。同時,收縮模塊通過跳轉(zhuǎn)注意力與跳轉(zhuǎn)連接向擴(kuò)張模塊傳遞多尺度信息。

圖2 分割網(wǎng)絡(luò)整體結(jié)構(gòu)Fig.2 The overall structure of segmented network

1.2.1 收縮模塊

收縮模塊由3個獨立的Swin Transformer結(jié)構(gòu)組成,對應(yīng)圖2中的編碼層。每一個Swin Transformer組成相同并且與傳統(tǒng)的結(jié)構(gòu)相同,如圖3所示,W-MSA表示窗口注意力,SW-MSA表示滑動窗口注意力,MLP表示多層感知機(jī)。Swin Transformer的第一層是歸一化層(LN),第二層是多頭自注意力層,第三層是第二個歸一化層,第四層是多層感知機(jī),對每個像素點分類。該模塊的計算步驟如下:

圖3 單個Swin Transformer模塊結(jié)構(gòu)Fig.3 The architecture of a Swin Transformer

(1)

(2)

(3)

(4)

式中:l為第l層;z為圖4中不同模塊的輸出。每一層的注意力計算公式具體如下:

圖4 訓(xùn)練過程Fig.4 The training process

Q=ΩqX+βqlT
K=ΩkX+βklT
V=ΩvX+βvlT

(5)

式中:Q、K和V分別為查詢向量、鍵向量和值向量;βq、βk、βv∈RD,βq、βk和βv為可選擇偏置項;Ωq∈Rd×dq,Ωk∈Rd×dk,Ωv∈Rd×dv;lT為偏置項的轉(zhuǎn)置。

圖像數(shù)據(jù)通過歸一化層后削弱噪聲影響并且減小數(shù)據(jù)間的差距,這使得傳入多頭自注意力層的數(shù)據(jù)能夠在分頭后每一部分都具有源數(shù)據(jù)的特征,進(jìn)而保證多層感知機(jī)輸出的特征是合理的,并且模型整體訓(xùn)練速度不會降低。每一個Swin Transformer塊輸出特征的尺寸是遞減的,維度是遞增的,這樣能夠確保特征圖在收縮部分和擴(kuò)張部分相呼應(yīng),進(jìn)而保證特征圖能夠恢復(fù)原尺寸。因為模型過多參數(shù)的更新和擴(kuò)張部分的反卷積操作通常會造成數(shù)據(jù)的丟失,出現(xiàn)無法恢復(fù)原有尺寸的問題,這些問題增加了推理時的數(shù)據(jù)處理的難度。

1.2.2 擴(kuò)張模塊

擴(kuò)張和收縮是對稱的,這是典型的U形結(jié)構(gòu)設(shè)計。擴(kuò)張模塊由3個結(jié)構(gòu)相同的Transformer塊組成,對應(yīng)圖2中的解碼層。Transformer塊的子模塊堆疊順序是相同,首先是歸一化層,其次是多頭自注意力層,第三層是第二個歸一化層,最終是多層感知機(jī)。收縮部分傳入的數(shù)據(jù)通過擴(kuò)張部分第一個Transformer塊的歸一化層后,特征的差距減小,這使得傳入多頭自注意力層的數(shù)據(jù)能夠在分頭后每一部分都具有近似的數(shù)學(xué)性質(zhì),再做一次歸一化依然可以起到這種作用。另外,還能夠加快參數(shù)更新速度,進(jìn)而減小模型的開銷,并且保證模型的整體分割能力不會下降。每一個Transformer塊輸出特征的尺寸是遞增的,維度是遞減的,這樣能夠確保特征圖在收縮部分傳來的特征可以輕松地與擴(kuò)張部分生成的特征合并,減少信息損失,提高了模型對數(shù)據(jù)的敏感度,有利于模型學(xué)習(xí)數(shù)據(jù)特征。

1.2.3 跳轉(zhuǎn)連接

跳轉(zhuǎn)連接的操作在收縮擴(kuò)張這類U形結(jié)構(gòu)中很常見,本文不僅使用了特征跳轉(zhuǎn)連接還使用了注意力跳轉(zhuǎn)操作。特征跳轉(zhuǎn)連接就是將收縮部分的特征直接傳入擴(kuò)張部分,直接進(jìn)行數(shù)學(xué)意義上的相加,相加后的新的特征作為擴(kuò)張部分的特征,這種操作是具有開創(chuàng)性意義的,經(jīng)過驗證是有效的,模型的分割效果也是有明顯提升的。Transformer在訓(xùn)練過程中會生成大量的數(shù)據(jù),這些數(shù)據(jù)的計算很考驗CPU的能力。為了解決這個問題,本文提出的模型將收縮部分計算出的K、V和Q值直接代替擴(kuò)張部分Transformer對應(yīng)層的對應(yīng)值,這種操作使模型在計算上花費(fèi)的時間較小,從而提升模型的訓(xùn)練速度。

2 實驗

2.1 評價指標(biāo)

為了更好地衡量模型的性能,采用了圖像分割領(lǐng)域常用的評估指標(biāo)Dice相似系數(shù)(dice similarity coefficient,DSC)來測量增強(qiáng)腫瘤區(qū)域(ET)、腫瘤核心區(qū)域(TC)以及整個腫瘤區(qū)域(WT)的分割精度。在腦膠質(zhì)瘤圖像分割中,Dice得分可以用來比較自動分割結(jié)果和人工分割結(jié)果(或真實標(biāo)注)之間的一致性,反映分割的精度和覆蓋度。豪斯多夫距離(hausdorff distance,HD)是一種用于衡量兩個集合之間的距離的指標(biāo),常用于評價圖像分割的精細(xì)度。

(6)

式中:TP為模型正確預(yù)測組織正樣本的數(shù)量;FP為模型錯誤預(yù)測組織正樣本的數(shù)量;FN為模型錯誤預(yù)測組織負(fù)樣本的數(shù)量。

2.2 實現(xiàn)細(xì)節(jié)

本文所有實驗都是基于Ubuntu16.04平臺進(jìn)行,模型的訓(xùn)練框架為Pytorch1.8,Python3.9。在訓(xùn)練模型時,將輸入圖像的尺寸縮放到128×128×128,學(xué)習(xí)率設(shè)置為1e-4,學(xué)習(xí)率調(diào)整策略設(shè)置為余弦衰減,優(yōu)化器設(shè)置為Adam優(yōu)化器,batch大小設(shè)為4,訓(xùn)練100個epoch,并采用包含32G內(nèi)存的Tesla V100顯卡對模型訓(xùn)練進(jìn)行加速。網(wǎng)絡(luò)的完整訓(xùn)練過程見圖4,可以發(fā)現(xiàn),隨著迭代次數(shù)的增加,模型的訓(xùn)練損失穩(wěn)步下降,直至收斂。

2.3 實驗結(jié)果與分析

為了驗證本文提出網(wǎng)絡(luò)結(jié)構(gòu)的有效性,在BraTS2021數(shù)據(jù)集上進(jìn)行了大量的對比實驗。實驗結(jié)果見表1,本文提出的網(wǎng)絡(luò)結(jié)構(gòu)在ET、TC和WT區(qū)域上的Dice得分分別達(dá)到了83.51%、87.66%以及91.39%。相比其他5種分割模型,其中ET區(qū)域的得分比最高的3D U-Net提升0.12%,比最低的UNTER提升3.73%;TC區(qū)域的得分比最高的nnFormer提升1.18%,比最低的UNTER提升4%;WT區(qū)域的得分比最高的nnFormer提升1.02%,比最低的TransBTS提升2.14%。在性能指標(biāo)上,本文提出模型的參數(shù)量只比3D U-Net高出8.9 M,但是浮點運(yùn)算量卻減少了392.9 GFlops。而相比TransBTS、UNTER、V-Net和nnFormer等模型,本模型的參數(shù)量分別減小了12.2、81.7、48.5和18.9 M。此外,本模型的浮點運(yùn)算次數(shù)只比nnFormer增加了54.3 G,參數(shù)量卻減小了近一倍。而相比TransBTS、UNTER、V-Net和3D U-Net模型,浮點運(yùn)算次數(shù)則分別減少了168.0、28.5、600.9和392.9 G。綜合整個實驗結(jié)果,本文提出的分割模型在精度和速度上取得了較好的平衡,更接近臨床應(yīng)用的目標(biāo)。

表1 實驗性能對比

2.4 消融實驗

為了探討網(wǎng)絡(luò)設(shè)計中重要參數(shù)對模型性能的影響,本節(jié)針對不同參數(shù)進(jìn)行了消融實驗。關(guān)于預(yù)訓(xùn)練權(quán)重對模型性能的影響結(jié)果如表2所示,在使用預(yù)訓(xùn)練權(quán)重后模型的整體性能都有所提升。其中,ET、TC、WT的Dice得分分別提高2.02%、1.34%、0.29%,可以發(fā)現(xiàn),使用預(yù)訓(xùn)練權(quán)重對于難以識別的小目標(biāo)提升更大,例如ET。

表2 預(yù)訓(xùn)練權(quán)值對模型的影響

關(guān)于跳轉(zhuǎn)注意力對模型性能的影響結(jié)果如表3所示,在使用跳轉(zhuǎn)注意力后,模型對ET和TC的Dice得分分別提高1.11%和1.39%,而WT的Dice得分僅提升0.13%。

表3 跳轉(zhuǎn)注意力對模型的影響

3 結(jié)論

本文針對卷積神經(jīng)網(wǎng)絡(luò)在提取圖像特征時存在的全局歸納偏差能力受限的問題,提出了一種完全基于Transformer的端到端的U形結(jié)構(gòu)網(wǎng)絡(luò),能夠直接處理三維的醫(yī)學(xué)圖像,并通過跳轉(zhuǎn)連接降低了Transformer在小數(shù)據(jù)集上的過擬合風(fēng)險。實驗結(jié)果表明,本文提出的網(wǎng)絡(luò)結(jié)構(gòu)相比TransBTS、UNTER、V-Net和3D U-Net等網(wǎng)絡(luò)具有更強(qiáng)的分割性能。

然而,本文在訓(xùn)練模型時,使用了Transformer在imagenet-1k數(shù)據(jù)集上的預(yù)訓(xùn)練權(quán)重對模型參數(shù)進(jìn)行初始化,模型的分割精度有所提升,這說明Transformer模型對于大數(shù)據(jù)集的依賴程度比較高。如果不使用跳轉(zhuǎn)注意力,不將模型收縮部分的K、V、Q值傳入擴(kuò)張部分,那么模型不僅會降低訓(xùn)練速度,分割性能也會降低。在下一步的研究工作當(dāng)中,將針對這些問題對模型做進(jìn)一步的優(yōu)化,提高模型的識別精度與速度,爭取達(dá)到臨床應(yīng)用的目標(biāo)。