TransFusionNet：JetsonTX2下肝腫瘤和血管分割的語義和空間特征融合框架

IEEE JOURNAL OF BIOMEDICAL AND HEALTH INFORMATICS, VOL. 27, NO. 3, MARCH 2023

背景

CT圖像中肝臟腫瘤和血管的分割和重建可以為醫生的術前計劃和手術干預提供便利。介紹了一個TransFusionNet框架，該框架由語義特征提取模塊、局部空間特征提取模塊，邊緣特征提取模塊和多尺度特征融合模塊組成，以實現肝臟腫瘤和血管的細粒度分割。此外，我們將遷移學習方法應用于使用公共數據集進行預訓練，然后對模型進行微調，以進一步提高擬合效果。此外，我們提出了一種智能量化方案來壓縮模型權重，并在JetsonTX2上實現了高性能推理。TransFusionNet框架在血管分割任務中實現了0.854的平均IoU，在肝腫瘤分割任務中獲得了0.927的平均IoU。在分析量化推理的計算性能時，我們的量化模型在NVIDIA RTX3090的Node上實現了4TFLOP，在JetsonTX2上實現了132GFLOP。這種前所未有的分割效果在一定程度上解決了自動分割的準確性和性能瓶頸。
需要解決的問題：
（1）如何設計一種準確、快速的肝腫瘤血管自動分割和三維重構方法？
（2） 如何設計一個能夠學習空間語義融合特征的分割框架，以提高腫瘤和血管細節的分割精度？
（3） 如何設計高效的模型量化方法，實現腫瘤和血管的高性能模型推斷和3D重建？
（4） 如何優化模型的計算和存儲開銷，以構建輕量級模型并將分割模型部署到JetsonTX2設備？

貢獻

1）我們提出了TransFusionNet框架，該框架結合了CT圖像的空間、語義和邊緣特征，實現了對肝臟腫瘤和肝內動脈血管系統的精確精細分割。
2） 我們提出了一種基于強化學習的智能量化方案來壓縮模型的權重，使模型在JetsonTX2和NVIDIA RTX3090 GPU的節點上都達到了最佳的推理性能。
3） 通過仔細量化，我們的模型使用NVIDIA RTX3090 GPU和JetsonTX2設備在Node上實現了高性能的肝血管腫瘤推斷和三維重建。

實驗

數據集：LITS（肝和肝腫瘤分割，https://competitions.codalab.org/competitions/17094)數據集包含130例腫瘤、轉移瘤和囊腫。這些CT掃描具有較大的空間分辨率和視場（FOV）差異[4]。
3Dircadb（用于算法數據庫比較的3D圖像重建，https://www.ircad.fr/research/3d-ircadb-01/)是一個公共數據集，可用于訓練和測試肝血管分割方法，包括20名不同圖像分辨率、血管結構、強度分布和肝血管對比CT增強的患者[29]。
同時，我們收集了18例患者的CT增強圖像，并構建了肝腫瘤血管（LTBV）數據集。為了清楚地區分肝臟中的血管并減輕標記負擔，我們只保留了18名患者的動脈期圖像。由于可訓練樣本的稀缺性，我們只分為訓練數據集和測試數據集。我們使用這兩個數據集進行模型預訓練，并使用我們的私人數據集對肝動脈和腫瘤分割任務的模型進行微調訓練。以上三個數據集按照8:2的比例分為訓練集和測試集。

消融實驗

方法

Transformer-Based Semantic Feature Extraction Module（transformer語義特征提取模塊）

有效解決了傳統深度CNN網絡中感知場缺陷導致的信息缺失問題
可以學習全局特征表示的編碼器，該編碼器由基于特征提取主干的特征嵌入模塊和基于Transformer感測圖像的語義相關信息表示的特征提取模塊組成[31]。該模塊采用了一種全新的特征提取思想，通過對圖片特征進行語義表示，并學習語義特征的全局表示。

Local Spatial Feature Extraction Module（局部空間特征提取模塊）

基于Transformer的特征提取模塊是一個非常強大的語義信息特征提取模塊，因為Transformer特征提取模塊在學習語義相關特征方面具有優勢。然而，在許多方面，Transformer并不是傳統卷積運算的有效替代品。對于提取一些圖像中更細微的特征，如感興趣區域的特征和微小血管的特征，CNN是一個完美的解決方案。我們設計了一種基于多層SEBottleNet堆疊的局部殘差網絡編碼器。我們在BottleNet中引入了擠壓和激勵（SE）[51]，以增強特征圖通道之間的相互依賴性。
確保每個模塊都有不同的特征提取任務。我們在模塊中間引入了擠壓和激勵模塊，以更好地學習特征圖通道維度的重要性，從而使SEBottleNet在特征提取過程中具有更強的學習重點。通過SEBottleNet和maxpool的連續疊加，編碼器可以連續提取輸入CT圖像的局部特征表示。同時，由于每個SEBottleNet都設置了剩余連接，這使編碼器能夠有效地緩解網絡深化帶來的降級問題。

Edge Extraction Module（邊緣提取模塊）

由于肝動脈血管非常小，因此進一步細化血管和肝臟的分割是一項具有挑戰性的任務。為了讓模型能夠學習更詳細的空間特征，我們引入了EEM，它專門用于學習感興趣的血管和腫瘤區域的邊緣特征，并將邊緣特征融合到分割網絡中。

EEM以特征提取層的特征圖和Canny算法[18]提取的CT邊緣圖（圖3（b））為輸入，預測邊緣結果e∈RH×W。該模塊預測邊緣信息，并將預測的特征圖組合到分割網絡中。為了完成這項任務，我們處理分割注釋以獲得邊緣注釋er（圖第3（d）段），它可以用作該模塊的監督條件。

在本模塊中，我們使用了門控激勵卷積（GEC）層。GEC是EEM中最重要的單元，它可以過濾掉一些不相關的信息，集中精力提取圖像的邊緣特征。GEC應用于EEM和特征提取模塊之間。它使用門控機制來停用其自身的激活，這些激活被提取模塊[52]中包含的高級信息認為不相關。同時，我們在門控激活層中引入了一個激勵模塊，以學習不同特征圖的重要性。

從理論上講，GEC可以簡單地看作是對特征圖的空間維度和通道維度的關注的集合。通過GEC運算，注意力圖αi選擇性地保留了邊緣語義特征。我們取消了對特征提取器的淺層特征圖的GEC操作，因為饋送到卷積層的圖像主要學習一般的低級特征，同時，輸出的特征圖保留了豐富的邊緣信息。隨著網絡的深入，特征圖將保留高級特征。在理論上，使用GEC運算可以有效地對高級特征的有用邊緣信息進行加權。

Canny算子可以有效地濾除圖像中不相關的特征，得到如圖所示的Canny圖像。第3（b）段。我們認為它適用于醫學圖像分割。因此，我們首先將canny圖像和最后一個GEC模塊輸出en連接起來。然后，我們將它們與兩個特征提取器的輸出特征圖一起提供給融合模塊。同時，邊緣提取模塊以邊緣損失為損失函數，以邊緣標簽為監督，對預測邊緣圖進行優化。

Multi-Scale Feature Fusing Module（多尺度特征融合模塊）

多尺度特征融合解碼模塊，以對三個模塊學習到的語義特征進行采樣。該模塊將三個模塊提取的特征圖作為輸入，并輸出預測的類別分布圖

主要融合三個特征提取模塊的特征圖。參考空間金字塔池（SPP）來設計模塊。首先，該模塊使用C1×1和C3×3卷積分別從語義特征圖和空間特征圖的拼接結果中提取特征。接下來，我們將其輸入到池化層，并融合邊緣特征圖。通過上述操作，獲得了三個不同感受野的特征圖。最后，我們對這三個特征圖進行采樣并連接，以輸出融合的特征圖。從理論上講，這些模塊輸出的特征圖可以保留豐富的空間特征、語義相關特征和邊緣特征。在連續融合不同尺度的低級特征圖后，解碼器可以從粗到細地學習圖像的語義信息。

Multi Task Training Strategy（多任務訓練策略）

分割任務、邊界任務、
優化過程中，特征提取模塊和邊緣提取模塊的參數將基于損失進行優化。接下來，我們將邊緣提取模塊輸出的特征圖輸入到融合模塊中，以預測分割結果。因此，邊緣提取模塊學習到的先驗知識保留在y中。同時，在優化過程中，分割損失會更加關注邊緣特征

Applying Transfer Learning to TransFusionNet（遷移學習在TransFusionNet上的應用）

由于癌癥篩查后增強CT圖像的稀缺性以及腫瘤和血管的標記困難，我們獲得了18例患者的CT圖像。數據過少將不可避免地影響模型的性能，并加深過度擬合問題。為此，我們引入了一種遷移學習策略，該策略不需要具有確切代表性的訓練數據，并且能夠利用數據集之間的相似性，在模型的訓練階段捕獲特定的先驗知識，以構建新的分割模型。

首先使用公共數據集LITS和3Dircadb對模型進行預訓練，分別獲得肝臟腫瘤分割模型和肝臟血管分割模型。然后，我們使用我們的肝臟腫瘤數據和肝臟血管數據來重新訓練通過預訓練獲得的模型。當我們需要對肝臟腫瘤和CT圖像的血管進行分割時，我們只需要輸入一張CT圖像，模型就會分別對CT圖像中的腫瘤和血管部分進行分割。

Quantification and Fine-Tuning of Inference Models（推理模型的量化與微調）

提出了一種基于硬件感知自動量化（HAQ）[55]的模型量化方案，以壓縮框架的CNN和密集層，并優化計算和存儲開銷

損失函數

分割損失

邊界損失

形狀正則化損失，y是預測的分割圖，e是預測的邊緣圖。在模型訓練開始時，由于邊緣提取模塊是隨機初始化的，無法準確預測e，（10）不起任何作用。因此，我們引入了一種動態調整策略，在100個歷元之前將λ4設置為0，在100歷元之后將λ4≥0

Thinking

雙編碼器分支，語義特征編碼器在中間層加入Transformer，SEBottleNet分支提取局部特征，融合模塊，邊界特征提取模塊，用標簽的canny邊緣圖作為監督