基于Transformer的圖像分割研究

譚棚文 向紅朵

（重慶師范大學，重慶 401331）

0 引言

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

該文的主要貢獻有以下2點：1） 在每層的跳躍連接過程中引入senet模塊進行處理后再進行拼接，以增強底層語義和高層語義的連接效果，從而為模型在進行多尺度預測和分割時提供更加精細的特征。2） 再引入一個基于Dice相似系數與交叉熵損失函數改進的損失函數，以解決訓練時可能出現的梯度消失情況。

1 研究現狀

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

2 模型設計

相關研究顯示，基于Transformer的相關方法具有更好的分割效果，該方法采用的Transformer模型具有“U”形編碼器-解碼器設計，可對輸入體數據進行整體處理，其網絡結構圖如圖1所示。

圖1 網絡結構圖

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

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

編碼器有2個連續的自注意層來同時編碼局部和全局信息，解碼器基于并行移動窗口的自注意模塊和交叉注意模塊來捕獲細節，同時接受同一階段編碼器生成的鍵（）與值（），并通過傅里葉位置編碼來細化邊界，在網絡跳躍連接時加入senet模塊，盡可能通過通道的特征去關注全局的目標特征。

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

式中：為Squeeze得到的結果；為Sigmoid函數；為ReLU 函數；，（）為全連接操作；F（，）為對壓縮結果進行自適應重新校正操作；為加權后的結果；（，）為對其進行ReLu激活；為數據矩陣；為一個縮放參數；為數據維度。

在經過Squeeze操作與Excitation操作后，在Reweight操作中，將學到的各個通道的激活值乘以原始特征得到最后的結果。具體操作如公式（3）所示。

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

3 試驗及結果分析

3.1 試驗環境與數據預處理

該文的試驗環境為CUDA V11.2并行計算架構，GPU型號是顯存為25.4 GB的RTX3090，采用Microsoft Windows 10操作系統與V1.10版本的開源PyTorch框架，運用Python3.7.4 在VSCode上對整體代碼進行編寫并運行。

該文的試驗數據來源于MICCAI BRATS2021腦膠質瘤數據集。該數據集由形狀為240·240·155的1 251次MRI影像組成，該文將這1 251次影像劃分為3組（800次、225次以及226次），分別用于訓練、驗證以及測試。

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

3.2 損失函數

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

式中：為輸出通道數；S為激活后神經網絡的輸出；R為基礎真值標簽；為平滑因子，在該文實驗中的數值為1；LOSS為Dice損失函數。

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

3.3 評價指標

該文采用Dice相似系數作為性能評價指標，其主要的評估指標是重疊度量和距離度量，具體定義如公式（6）所示。

式中：為真陽性；為假陽性；為假陰性。

Dice系數的范圍為0～1，當Dice系數為1時，表示完全重疊。

3.4 模型訓練與驗證

該文采用了Liu Ze等人在ImageNet-1K圖像分類中預先訓練的權重來初始化模型。訓練時采用了Adam優化算法，網絡批大小設置為1，訓練周期設置為200，訓練時初始學習速率為0.000 1，采用余弦衰減速率調度器重新設置學習速率，衰減步數為每4個周期衰減1次。模型是通過五重交叉驗證產生的，驗證集僅在訓練期間用于調整超參數，并在訓練過程結束時對其性能進行基準測試。在模型訓練過程中，損失函數的變化曲線如圖2所示。由圖2可知，隨著迭代次數的增加，損失值在逐漸降低，當訓練迭代次數大約為170時，網絡損失已趨于穩定。

圖2 損失函數變化情況圖

3.5 實驗結果

每個腫瘤類別在訓練周期時的得分情況如圖3所示。

圖3 Dice評分收斂圖

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

圖4 原始數據與5種模型的分割結果圖

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

3.6 消融實驗

為進一步研究各方法對整體性能的貢獻，該文逐步將不同的組件集成到模型中去。由表1～表4可知，所有模塊都有助于提高模型的性能，使用預先訓練的權重在一定程度上可以提高模型的分割性能，與使用單個損失函數相比，混合損失函數的應用在一定程度上解決了網絡訓練過程中梯度消失的情況。同時，引入senet模塊進一步加強了模型的分割性能。

表1 5種模型的分割性能對比表

表4 senet模塊的消融研究

表2 預訓練模型的消融研究

表3 損失函數的消融研究

4 結語

準確地對腦膠質瘤進行分割對腦癌的早期診斷非常重要。在將注意力機制應用于模型中的跳躍連接過程中，采用改進的混合函數來優化模型，使其能夠更好地幫助模型學習樣本。該文的研究結果表明，與現有的傳統分割方法相比，所做出的改進在腦膠質瘤分割方面有一定提升。然而該文所提出的算法具有一定的局限性，當訓練數據集較少時，所得到的結果會較低。因此使用有限樣本得到較好分割性能是下一步研究的重點。