基于Transformer的圖像分割研究

譚棚文 向紅朵

（重慶師范大學(xué)，重慶 401331）

0 引言

腦膠質(zhì)瘤是對(duì)神經(jīng)系統(tǒng)膠質(zhì)細(xì)胞腫瘤和神經(jīng)元細(xì)胞腫瘤的統(tǒng)稱，其發(fā)病率高，治愈率低。目前，對(duì)該疾病的診斷以及對(duì)手術(shù)后病人恢復(fù)效果的評(píng)估主要依賴于磁共振成像（Magnetic Resonance Imaging，MRI）技術(shù)，而腦膠質(zhì)瘤的分割往往依靠醫(yī)生根據(jù)臨床經(jīng)驗(yàn)進(jìn)行手動(dòng)分割，該方法不僅會(huì)耗費(fèi)大量時(shí)間與精力，而且分割準(zhǔn)確率難以得到保障。因此，設(shè)計(jì)一個(gè)采用自動(dòng)或半自動(dòng)方法且能夠精準(zhǔn)分割腦膠質(zhì)瘤的方法十分重要，這不僅可以提供精確的治療方案，而且還可以優(yōu)化患者的預(yù)后。

該文的主要貢獻(xiàn)有以下2點(diǎn)：1） 在每層的跳躍連接過(guò)程中引入senet模塊進(jìn)行處理后再進(jìn)行拼接，以增強(qiáng)底層語(yǔ)義和高層語(yǔ)義的連接效果，從而為模型在進(jìn)行多尺度預(yù)測(cè)和分割時(shí)提供更加精細(xì)的特征。2） 再引入一個(gè)基于Dice相似系數(shù)與交叉熵?fù)p失函數(shù)改進(jìn)的損失函數(shù)，以解決訓(xùn)練時(shí)可能出現(xiàn)的梯度消失情況。

1 研究現(xiàn)狀

目前，卷積神經(jīng)網(wǎng)絡(luò)（Convolutional Neural Networks，CNN）已經(jīng)廣泛應(yīng)用于醫(yī)學(xué)影響分割領(lǐng)域，并表現(xiàn)出較為優(yōu)異的性能。例如TransUNet模型將Transformer作為完成醫(yī)學(xué)圖像分割任務(wù)的強(qiáng)大編碼器，并借助UNet恢復(fù)localized空間信息，然而其無(wú)法在降低模型參數(shù)量的同時(shí)保持高分割精度。類似Unet的純Transformer網(wǎng)絡(luò)模型Swin-Unet，使用帶有偏移窗口的分層Swin Transformer作為編碼器來(lái)提取上下文特征，然而該方法無(wú)法充分利用圖像的三維特征。3DUnet的誕生為醫(yī)學(xué)影像分割提供了極大的幫助，在很大程度上解決了3D圖像分割時(shí)需要將一個(gè)個(gè)切片輸入模型進(jìn)行訓(xùn)練的復(fù)雜問(wèn)題，也大幅提升了訓(xùn)練效率。HATAMIZADEH A等人提出了UNETR網(wǎng)絡(luò)模型，將三維數(shù)據(jù)分割任務(wù)重新設(shè)計(jì)為序列到序列的預(yù)測(cè)問(wèn)題，以提高圖像分割性能，但在分割過(guò)程中容易缺失圖像的局部信息。還有一種Gated Axial-Attention模型，通過(guò)在自注意力模塊中引入附加的控制機(jī)制來(lái)擴(kuò)展現(xiàn)有體系結(jié)構(gòu)，然而其在訓(xùn)練過(guò)程中引入了過(guò)多的冗余信息，導(dǎo)致其分割精度較低。

2 模型設(shè)計(jì)

相關(guān)研究顯示，基于Transformer的相關(guān)方法具有更好的分割效果，該方法采用的Transformer模型具有“U”形編碼器-解碼器設(shè)計(jì)，可對(duì)輸入體數(shù)據(jù)進(jìn)行整體處理，其網(wǎng)絡(luò)結(jié)構(gòu)圖如圖1所示。

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

在senet中，比較重要的就是壓縮操作與自適應(yīng)重新校正操作，與只在一個(gè)局部空間進(jìn)行操作而無(wú)法獲取足夠信息的傳統(tǒng)卷積相比，senet 設(shè)計(jì)了Squeeze操作，其每個(gè)通道的具體操作如公式（1）所示。

式中：f(u)為對(duì)u矩陣進(jìn)行Squeeze 操作；為高度；為寬度；u為矩陣中的第行、第列元素；Z為壓縮矩陣中的第個(gè)元素。

編碼器有2個(gè)連續(xù)的自注意層來(lái)同時(shí)編碼局部和全局信息，解碼器基于并行移動(dòng)窗口的自注意模塊和交叉注意模塊來(lái)捕獲細(xì)節(jié)，同時(shí)接受同一階段編碼器生成的鍵（）與值（），并通過(guò)傅里葉位置編碼來(lái)細(xì)化邊界，在網(wǎng)絡(luò)跳躍連接時(shí)加入senet模塊，盡可能通過(guò)通道的特征去關(guān)注全局的目標(biāo)特征。

Squeeze操作將原始維度為（其中，為高度；為寬度，；為通道數(shù)）的數(shù)據(jù)壓縮為11，也就是將三維數(shù)據(jù)壓縮成一維數(shù)據(jù)。同時(shí)，為了利用通道間的相關(guān)性來(lái)訓(xùn)練真正的目標(biāo)維數(shù)，為了限制模型的復(fù)雜度并增強(qiáng)其泛化能力，自適應(yīng)重新校正部分在門限機(jī)制中使用bottleneck形式的2個(gè)全連接層，具體操作如公式（2）所示。

式中：為Squeeze得到的結(jié)果；為Sigmoid函數(shù)；為ReLU 函數(shù)；，（）為全連接操作；F（，）為對(duì)壓縮結(jié)果進(jìn)行自適應(yīng)重新校正操作；為加權(quán)后的結(jié)果；（，）為對(duì)其進(jìn)行ReLu激活；為數(shù)據(jù)矩陣；為一個(gè)縮放參數(shù)；為數(shù)據(jù)維度。

在經(jīng)過(guò)Squeeze操作與Excitation操作后，在Reweight操作中，將學(xué)到的各個(gè)通道的激活值乘以原始特征得到最后的結(jié)果。具體操作如公式（3）所示。

式中：u為1個(gè)二維矩陣；F（u,s）u與s進(jìn)行逐通道相乘；X為升維后的矩陣中的第個(gè)元素。

3 試驗(yàn)及結(jié)果分析

3.1 試驗(yàn)環(huán)境與數(shù)據(jù)預(yù)處理

該文的試驗(yàn)環(huán)境為CUDA V11.2并行計(jì)算架構(gòu)，GPU型號(hào)是顯存為25.4 GB的RTX3090，采用Microsoft Windows 10操作系統(tǒng)與V1.10版本的開源PyTorch框架，運(yùn)用Python3.7.4 在VSCode上對(duì)整體代碼進(jìn)行編寫并運(yùn)行。

該文的試驗(yàn)數(shù)據(jù)來(lái)源于MICCAI BRATS2021腦膠質(zhì)瘤數(shù)據(jù)集。該數(shù)據(jù)集由形狀為240·240·155的1 251次MRI影像組成，該文將這1 251次影像劃分為3組（800次、225次以及226次），分別用于訓(xùn)練、驗(yàn)證以及測(cè)試。

該文的輸入數(shù)據(jù)為四模態(tài)MRI影像，其影像數(shù)據(jù)會(huì)因采集設(shè)備的制造商、各影像的采集參數(shù)和序列的不同而產(chǎn)生差異。為了使網(wǎng)絡(luò)更好地工作，需要對(duì)輸入圖像進(jìn)行標(biāo)準(zhǔn)化處理。首先，對(duì)所有MRI序列影像進(jìn)行線性歸一化（Min-Max Scale）。其次，使用包括整個(gè)大腦的最小邊界框，并通過(guò)移除不必要的背景將體積裁剪成128·128·128的固定大小，這樣可以去除原始影像中大部分無(wú)用的背景，并保留有效影像。最后，為了防止過(guò)擬合，該文通過(guò)多種方法按照一定的概率進(jìn)行動(dòng)態(tài)數(shù)據(jù)增強(qiáng)。具體方法以及各自的概率如下：沿每個(gè)空間軸旋轉(zhuǎn)（概率為60%）。加入高斯噪聲，使用標(biāo)準(zhǔn)偏差為0.1的中心正態(tài)分布；輸入通道重新縮放，將每個(gè)體素乘以0.9和1.1之間均勻采樣的因子（概率為80%）。

3.2 損失函數(shù)

由于Dice損失函數(shù)（）測(cè)量的是2個(gè)輪廓邊緣的最大距離，如果單個(gè)體素遠(yuǎn)離參考分割，那么Hausdorff距離值將會(huì)偏高。為了解決這個(gè)問(wèn)題，受Jiang Zeyu 等人在2019年Brats競(jìng)賽中解決方案的啟發(fā)，該文決定使用進(jìn)行訓(xùn)練，該損失函數(shù)不需要在前景和背景體素之間建立正確的平衡，也就是不用給不同類別的樣本分配權(quán)重。其損失值是按批次和按通道計(jì)算的，沒(méi)有加權(quán)，具體操作如公式（4）所示。

式中：為輸出通道數(shù)；S為激活后神經(jīng)網(wǎng)絡(luò)的輸出；R為基礎(chǔ)真值標(biāo)簽；為平滑因子，在該文實(shí)驗(yàn)中的數(shù)值為1；LOSS為Dice損失函數(shù)。

最后混合損失函數(shù)的定義如公式（5）所示。

3.3 評(píng)價(jià)指標(biāo)

該文采用Dice相似系數(shù)作為性能評(píng)價(jià)指標(biāo)，其主要的評(píng)估指標(biāo)是重疊度量和距離度量，具體定義如公式（6）所示。

式中：為真陽(yáng)性；為假陽(yáng)性；為假陰性。

Dice系數(shù)的范圍為0～1，當(dāng)Dice系數(shù)為1時(shí)，表示完全重疊。

3.4 模型訓(xùn)練與驗(yàn)證

該文采用了Liu Ze等人在ImageNet-1K圖像分類中預(yù)先訓(xùn)練的權(quán)重來(lái)初始化模型。訓(xùn)練時(shí)采用了Adam優(yōu)化算法，網(wǎng)絡(luò)批大小設(shè)置為1，訓(xùn)練周期設(shè)置為200，訓(xùn)練時(shí)初始學(xué)習(xí)速率為0.000 1，采用余弦衰減速率調(diào)度器重新設(shè)置學(xué)習(xí)速率，衰減步數(shù)為每4個(gè)周期衰減1次。模型是通過(guò)五重交叉驗(yàn)證產(chǎn)生的，驗(yàn)證集僅在訓(xùn)練期間用于調(diào)整超參數(shù)，并在訓(xùn)練過(guò)程結(jié)束時(shí)對(duì)其性能進(jìn)行基準(zhǔn)測(cè)試。在模型訓(xùn)練過(guò)程中，損失函數(shù)的變化曲線如圖2所示。由圖2可知，隨著迭代次數(shù)的增加，損失值在逐漸降低，當(dāng)訓(xùn)練迭代次數(shù)大約為170時(shí)，網(wǎng)絡(luò)損失已趨于穩(wěn)定。

圖2 損失函數(shù)變化情況圖

3.5 實(shí)驗(yàn)結(jié)果

每個(gè)腫瘤類別在訓(xùn)練周期時(shí)的得分情況如圖3所示。

圖3 Dice評(píng)分收斂圖

在圖4中，該文選取最近提出、效果較好且采用三維體數(shù)據(jù)作為輸入數(shù)據(jù)的其他模型并進(jìn)行對(duì)比。由圖4可知，該文提出的方法與原始數(shù)據(jù)中的分割影響相近，具有良好的分割性能。

圖4 原始數(shù)據(jù)與5種模型的分割結(jié)果圖

表1將該文的方法與基于Unet的3D Unet（模型三）方法、EquiUnet（模型一）和Attention EquiUnet（模型四）方法以及TrasnBTS（模型二）方法進(jìn)行對(duì)比。由表1可知，該文所提出的方法在分割腦膠質(zhì)瘤上的表現(xiàn)均好于其他網(wǎng)絡(luò)。

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

為進(jìn)一步研究各方法對(duì)整體性能的貢獻(xiàn)，該文逐步將不同的組件集成到模型中去。由表1～表4可知，所有模塊都有助于提高模型的性能，使用預(yù)先訓(xùn)練的權(quán)重在一定程度上可以提高模型的分割性能，與使用單個(gè)損失函數(shù)相比，混合損失函數(shù)的應(yīng)用在一定程度上解決了網(wǎng)絡(luò)訓(xùn)練過(guò)程中梯度消失的情況。同時(shí)，引入senet模塊進(jìn)一步加強(qiáng)了模型的分割性能。

表1 5種模型的分割性能對(duì)比表

表4 senet模塊的消融研究

表2 預(yù)訓(xùn)練模型的消融研究

表3 損失函數(shù)的消融研究

4 結(jié)語(yǔ)

準(zhǔn)確地對(duì)腦膠質(zhì)瘤進(jìn)行分割對(duì)腦癌的早期診斷非常重要。在將注意力機(jī)制應(yīng)用于模型中的跳躍連接過(guò)程中，采用改進(jìn)的混合函數(shù)來(lái)優(yōu)化模型，使其能夠更好地幫助模型學(xué)習(xí)樣本。該文的研究結(jié)果表明，與現(xiàn)有的傳統(tǒng)分割方法相比，所做出的改進(jìn)在腦膠質(zhì)瘤分割方面有一定提升。然而該文所提出的算法具有一定的局限性，當(dāng)訓(xùn)練數(shù)據(jù)集較少時(shí)，所得到的結(jié)果會(huì)較低。因此使用有限樣本得到較好分割性能是下一步研究的重點(diǎn)。