Trans-SegNet:一種基于Transformer的腦腫瘤圖像分割網絡

仇龍

（北京科技大學，北京 100083）

腦腫瘤是一種常見的神經系統疾病，是指在腦部或其周圍組織中發生的一類惡性或良性腫瘤，其分布范圍廣泛，易于侵襲周圍組織，對患者的生命和健康造成威脅。全球每年新發腦腫瘤約25萬例。據統計，腦腫瘤在惡性腫瘤中占比不到2%，但是近年來其發病率持續上升，已成為21世紀危害人類健康的重要疾病之一。早期診斷和治療是腦腫瘤預防和根治的基礎[1]，但醫生在診斷中會受到自身經驗、知識水平、情感因素等的影響，導致診斷結果存在主觀性，容易出現誤診和漏診等問題。此外，由于醫生的不同背景、經驗和知識水平的不同，不同醫生對同一病例的診斷結果可能缺乏一致性。而醫學圖像分割技術是一種基于計算機算法的自動化技術，可以通過降低醫生分析的主觀因素，提高疾病診斷的客觀性和準確性。還可以幫助醫生更準確地評估腫瘤的位置、形狀和大小等信息，通過對圖像進行分割，可以清晰地顯示出腫瘤和周圍組織之間的分界線，減少誤診和漏診的風險。對于需要手術治療的患者，醫學圖像分割技術可以幫助醫生更準確地評估腫瘤的位置和大小，規劃手術方案，減少手術風險。

隨著深度學習技術的不斷發展，醫學圖像分割中基于卷積神經網絡的方法已經成為主流[2-3]，其中，UNet[4]、V-Net[5]和U-Net++等方法得到廣泛應用。UNet 基于卷積神經網絡采用了編碼器-解碼器結構。在編碼器部分，卷積和池化操作可用于降低輸入圖像的分辨率并提取圖像特征。在解碼器部分，卷積和上采樣操作被用來將特征圖恢復到原始圖像的分辨率，并進行像素級別的分類。此外，還引入了跳躍連接，將編碼器的特征圖送入解碼器中幫助解決分割任務中的信息丟失問題。雖然已經取得了較好的分割效果，但其由于卷積操作引入的固有偏置原因，仍存在無法建模長距離依賴和捕捉全局上下文信息的問題，這限制了分割精度和輔助診斷準確率的進一步提升[6-7]。近年來基于自注意力機制的神經網絡結構Transformer[8]在自然語言處理領域和圖像領域都取得了很好的效果，相比于傳統的循環神經網絡，Transformer 能夠處理更長的文本序列，并且能夠捕捉更復雜的語言關系。這使得Transformer 成了自然語言處理領域的一種重要技術。其通過內部的自注意力機制可以自動對圖像中不同位置的特征進行加權和組合，從而提高模型的表達能力并建立全局層面的長距離依賴。這種特性在醫學圖像分割中非常重要，因為醫學圖像中存在很多微小的細節和不規則的形狀，需要模型具有較強的表達能力才能更好地分割[9]。此外，全局信息可以幫助模型更好地理解圖像中的結構和背景，在腦部MRI 圖像中，腫瘤結構往往與大腦組織的結構緊密相關，需要考慮整個腦部結構才能更好地分割腫瘤。綜上，本文從腦腫瘤圖像分割算法中引入Transformer結構并利用自注意力機制，設計了新的醫學圖像分割網絡Trans-SegNet。實驗結果表明，Trans-SegNet可以有效地提升模型的建模能力并建立長距離全局信息，提升腦腫瘤分割的準確性。

1 Trans-SegNet分割算法設計

1.1 算法整體結構

本文基于傳統的U-Net 架構和Transformer 結構，提出了一個新的腦腫瘤圖像分割算法，其整體框架如圖1 所示。相較于原本U-Net，本文將Transformer 層多尺度并行插入原始U-Net 的跳躍連接中并設計全局融合模塊，補充不同尺度特征中的長距離依賴，并送入解碼器進一步融合。具體來說，所設計模型Trans-SegNet 包含3個主要部分，分別是編碼器、解碼器以及全局特征融合模塊。其中，編碼器由5個卷積塊組成，每個塊由3×3卷積、批歸一化和ReLU激活函數組成，其結構如圖1所示。并在每個卷積塊間逐步下采樣壓縮圖像尺寸，依次為原始圖像分辨率的1/4、1/8、1/16和1/32。同時，加入殘差連接的操作，使模型優化更快避免出現訓練不穩定的現象。得到多尺度特征后，由于Transformer 是處理序列信息的，需要首先對特征圖進行特征形狀變換，將3維的圖像特征變換為1維。為保留圖像塊之間的相對位置信息，在序列上附加位置編碼后，分別送入4個并行的全局特征融合模塊來同時建立4個尺度上的長距離特征關系。挖掘并增強腫瘤相似特征間的相關性，提取上下文語義信息，幫助模型更好地理解圖像中的結構和關系，區分腫瘤和正常組織，提升分割精度的準確性。將經過全局增強模塊后的特征圖進行形狀變換，恢復其3維結構，使其適配3維卷積的形狀結構后，分別送入解碼器中。解碼器中包含與編碼器類似的卷積塊，尺寸最小的特征圖通過卷積模塊和上采樣操作與經過全局增強后的多尺度特征拼接融合并逐階段恢復和生成用于最終分割的特征，實現腦腫瘤圖像的精準分割，下面將詳細介紹全局融合模塊。

圖1 Trans-SegNet模型結構圖

1.2 全局特征融合模塊結構

全局特征融合模塊由自注意力機制、頻域表示模塊、層歸一化和前向傳播層等組件構成。圖像送入全局特征融合模塊前在一個圖像塊層面被序列化，將3維的醫學圖像數據轉化成1維的序列數據，并附加位置編碼保持圖像相對位置關系。之后并行通過自注意力機制和頻域表示學習模塊建立整張圖像上的相互作用并捕捉上下文信息。在自注意力機制中，對于每個圖像塊，會將該圖像塊的特征表示作為查詢（Query）、鍵（Key）和值（Value），然后計算該圖像塊與所有其他該圖像塊之間的相似度。將所計算的相似度作為權重，對所有值進行加權求和，得到的加權和為該查詢的輸出[7]。其數學表達式如下：

式中,WQ,WK,WV表示可學習的參數，dk表示特征向量維度，x表示輸入特征。通過上述先計算全局圖像塊的關聯度，得到關聯親密度矩陣，再增強關聯性強的圖像塊，抑制關聯性弱的圖像塊的方法，模型的全局建模能力得到增強，圖像特征之間的相互作用被構建。

為進一步完善和補充圖像全局信息，本文設計了并行的頻域表示學習模塊。由于頻域可以反映圖像中的全局信息，其提供了對圖像中所有頻率的解析能力，可以通過分析頻域中的頻率，了解圖像中存在的所有周期性變化。包括圖像中的低頻分量（代表圖像中的大型特征和整體結構）和高頻分量（代表圖像中的細節和紋理）[9]，因此在頻域中分析是個合理可行的方案。具體來說，本文首先將特征經過傅里葉變換映射到頻域中，之后通過添加2個可學習的參數增強模型在頻域中豐富模型的全局特征表示能力，然后進行傅里葉反變換，與自注意力分支特征對齊，其數學表達式如下所示。

式中，FDB表示頻域表示學習模塊，IFFT、FFT分別表示傅里葉變換和反變換，a,b表示可學習參數，其形狀大小與輸入特征x相同，分別進行逐元素相乘和相加，通過這種方式可以有效地實現在頻域中學習全局信息。

特征經過全局信息增強后被送入前向傳播層中，其中包括兩個全連接層和一個激活層，激活層引入了非線性映射，補足自注意力機制無法處理非線性關系的問題。全連接層可以進一步增強模型的表示能力。其數學表達式如下：

式中,W1,W2,b1,b2表示可學習的參數，此外，全局特征融合模塊還引入了殘差連接和層歸一化等技術，來加速模型的訓練和提高模型的泛化能力。

綜上，全局特征融合模塊通過自注意力機制和頻域表示模塊捕獲增強了全局上下文信息，并從頻域角度完善特征間相互依賴。此外，還加入了前向傳播層以提升模型表示能力，實驗證明，通過以上改進可以顯著提升模型的分割精度。

2 實驗與結果

2.1 實驗細節和評價指標

本文基于PyTorch 框架并在2 張NVIDIA Geforce RTX3090 上訓練，顯存為24GB，使用PyTorch 版本為1.8.0，使用版本為11.3 的CUDA 和8.2 的cuDNN 進行訓練加速。實驗數據集為公開數據集BraTS2019[10]，它包括335個用于訓練的患者病例和125個用于驗證的病例。每個樣本都由具有4 種模式的3D 大腦MRI掃描組成。每個模態都有240×240×155的體積，已經對齊到同一空間，共包括背景、腫瘤全部區域(WT)、腫瘤核(TC)和增強腫瘤(ET)4 個類別。輸入經過裁剪后的圖片尺寸為128×128×128 以對齊空間分辨率和降低計算復雜程度。為避免過擬合，在實驗中增加隨機翻轉、調節對比度和隨機縮放以生成更多訓練樣本，并在測試時使用多尺度增強測試來提升模型性能。本文在通過五折交叉驗證，評估得分訓練樣本數為269，驗證樣本數為66，批處理大小4，學習率1e-4，訓練輪數500輪。本文以Dice得分來評估模型性能，其具體數學計算方式如下式所示。

式中A,B分別表示該類別圖像上的預測區域和真實標簽區域。此外，由于醫學腦腫瘤圖像存在前后景比例懸殊大，相對結構較為固定并且不同類別數量極度不平衡，本文不采用傳統圖像分割上的交叉熵損失，訓練階段優化的目標函數使用Dice 損失。Dice損失是基于相似性的評價指標，更看重重疊的比例。可以很好地處理類別像素值不平衡的問題。優化方面，由于傳統隨機梯度下降，可能會存在局部最優解和收斂速度慢的問題，本文選用了adam 的優化方法提升模型收斂速度。

2.2 實驗結果

本文在BraTS2019 上進行實驗，并驗證了所設計模型和具體模塊的有效性，實驗結果如表1所示。在相同實驗條件下，Trans-SegNet 相較于傳統基于卷積神經網絡的U-Net 模型在三種評價指標下均取得了大幅提升，其中ET 提升2.07%、WT 提升0.98%、TC 提升1.89%，證明了本文方法的優勢。具體來說，從模型角度分析，為了排除是骨干網絡的變化導致的精度差異，筆者設計了當Trans-SegNet 不添加全局特征模塊僅采用卷積神經網絡的實驗，結果相較于之前方法并沒有取得突出的性能優勢。為了進一步地驗證Transformer 插入到模型中的效果，筆者增加全局特征模塊但不包含頻域特征學習模塊的實驗，分割精度取得較大幅度提升1.31%(ET)、0.75%(WT)、0.98%(TC)，實驗結果證明，Transformer 的全局信息建模能力和多尺度插入融合方法，能夠有效地幫助模型理解上下文語義，提升分割準確率。此外，為驗證頻域表示學習模塊的有效性，筆者進行了添加頻域表示學習模塊的實驗即最終的Trans-SegNet，結果表明通過頻域表示學習模塊，模型能在頻域有效地完善全局信息，并進一步提升模型分割精度，實現腦腫瘤的精準分割。

表1 實驗結果精度對比

3 結論

針對目前深度學習算法在腦腫瘤圖像分割中存在全局特征和上下文信息提取能力弱以及分割精度低等問題，本文基于傳統的卷積神經網絡U-Net 算法提出了改進，本文通過在跳躍連接中添加全局特征增強模塊，輔助模型捕捉全局語義信息。具體來說，其中包括由頻域表示學習和自注意力機制組成的雙分支全局信息增強部分和前向傳播層以增強模型表示能力。實驗結果表明，相比改進前的傳統卷積神經網絡U-Net分割精度分別在ET/WT/TC上提升了2.07%、0.98%、1.89%。并進一步從實驗上證明了所設計頻域表示學習模塊和全局表示學習模塊的有效性以及Transformer 和在跳躍連接中插入的方式能夠有效地幫助模型具備全局建模能力，這有助于未來對腦腫瘤圖像分割的研究，給醫學圖像輔助診斷技術提供借鑒。