UNETR: Transformers for 3D Medical Image Segmentation
UNETR：用于三維醫(yī)學圖像分割的Transformer

Published: Oct 2021
Published in: IEEE Winter Conference on Applications of Computer Vision (WACV) 2022
論文：https://arxiv.org/abs/2103.10504
代碼：https://monai.io/research/unetr

摘要：
??過去十年以來，具有收縮路徑和擴展路徑（編碼器和解碼器）的全卷積神經(jīng)網(wǎng)絡（FCNNs）在各種醫(yī)學圖像分割應用中表現(xiàn)突出。在這些體系結構中，編碼器通過學習全局和局部特征以及上下文表示成為一個不可或缺的角色，這些特征和上下文表示可以被解碼器用于語義輸出預測。盡管它們?nèi)〉昧顺晒Γ?strong>但卷積層在FCNN中的局限性，限制了在此類網(wǎng)絡中學習遠程空間相關性的能力。受自然語言處理（NLP）在遠程序列學習中的最新成功的啟發(fā)（Transformers），作者將體積（3D）醫(yī)學圖像分割的任務重新設計為序列到序列的預測問題。介紹了一種稱為UNEt TRansformers（UNETR）的新穎架構，該架構利用純Transformers作為編碼器來學習輸入量的序列表示并有效地捕獲全局多尺度信息。同時也遵循了編碼器和解碼器的成功的“U型”網(wǎng)絡設計,Transformers編碼器通過不同分辨率的跳躍連接直接連接到解碼器，以計算最終的語義分段輸出。使用多器官分割的多圖譜標記(BTCV)數(shù)據(jù)集、醫(yī)學分割十項全能（MSD）數(shù)據(jù)集廣泛驗證了提出的模型在不同成像方式（即MR和CT）上對體積腦腫瘤和脾臟分割任務的性能，并且結果始終證明了良好的性能。

問題動機：
??圖像分割在定量醫(yī)學圖像分析中起著不可或缺的作用，因為它往往是解剖結構分析的第一步。自從深度學習出現(xiàn)以來，F(xiàn)CNNS，特別是“U型”編解碼架構已經(jīng)在各種醫(yī)學語義分割任務中取得了最先進的結果。在一個典型的U-Net架構中，編碼器負責通過逐步降低提取特征的采樣來學習全局上下文表示，而解碼器負責將提取的表示采樣到輸入分辨率，以進行像素/體素的語義預測。此外，跳過連接合并編碼器的輸出和解碼器在不同的分辨率，因此允許恢復在降采樣期間丟失的空間信息。
??雖然這種基于FCN的方法具有強大的表示學習能力，但它們在遠程依賴學習中的表現(xiàn)，僅限于它們的局部接受域。因此，這種在多尺度信息捕獲方面的缺陷導致了對不同形狀和尺度的結構(如不同大小的腦損傷)的次優(yōu)分割存在不足。一些研究已經(jīng)使用atrous convolutional layers來擴大感受域。然而，卷積層中接受域的局部性仍然限制了它們的學習能力到相對較小的區(qū)域。將自注意模塊與卷積層相結合，提高了非局部建模能力。
??在自然語言處理(NLP)中，基于Transformers的模型在各種任務中實現(xiàn)了最先進的基準。Transformers的自注意機制允許動態(tài)突出詞序列的重要特征。此外，在計算機視覺中，使用Transformers作為主干編碼器是有益的，因為它們具有建模遠程依賴和捕獲全局上下文的強大能力。具體來說，與卷積的局部公式不同，Transformers將圖像編碼為1維 patch嵌入序列，并利用自注意模塊學習從隱藏層計算出的值的加權和。因此，這種靈活的提法可以有效地學習長距離信息。此外，視覺Transformers (ViT)及其變體在學習文本前任務方面表現(xiàn)出了出色的能力，這些任務可以轉(zhuǎn)移到下游應用任務。
??在這項工作中，作者提出利用Transformers的能力進行體積醫(yī)學圖像分割，并引入了一種新的結構稱為UNEt TRansformers (UNETR)。特別的，作者將三維分割任務重新表述為一維序列對序列的預測問題，并使用Transformers作為編碼器從嵌入的輸入補丁中學習上下文信息。從Transformers編碼器中提取的表示通過多個分辨率的跳過連接與基于CNN的解碼器合并，以預測分割輸出。提出的框架使用基于CNN的解碼器，而不是在解碼器中也使用Transformers。這是由于盡管Transformers具有很強的全局信息學習能力，但卻不能正確地捕獲局部信息。
??作者使用Beyond the Cranial Vault (BTCV)和Medical segmentation Decathlon (MSD)數(shù)據(jù)集驗證了方法在3D CT和MRI分割任務中的有效性。在BTCV數(shù)據(jù)集中，UNETR在其排行榜上的標準和自由競賽部分都實現(xiàn)了新的最先進的表現(xiàn)。在MSD數(shù)據(jù)集中，UNETR在腦腫瘤和脾臟分割任務方面的表現(xiàn)優(yōu)于最先進的方法。

主要貢獻如下:
??1、提出了一種新的基于Transformers的醫(yī)學圖像分割模型。
??2、提出了一種新穎的架構，其中:(1) Transformers編碼器直接利用嵌入的3D體塊有效地捕捉遠程依賴;(2)跳躍式連接解碼器結合提取的不同分辨率的表示，并預測分割輸出。
??3、在兩個公共數(shù)據(jù)集(BTCV和MSD)上驗證了提出的模型對不同體積分割任務的有效性。UNETR在BTCV數(shù)據(jù)集排行榜上實現(xiàn)了新的最先進的性能，并在MSD數(shù)據(jù)集上優(yōu)于競爭方法。

思路來源：
??CNN-based分割網(wǎng)絡：自從引入開創(chuàng)性的U-Net以來，基于CNN的網(wǎng)絡已經(jīng)在各種2D和3D各種醫(yī)學圖像分割任務上取得了最先進的結果。對于以體積為單位的分割，有時使用三平面結構來組合每個體素的三視圖切片，也稱為2.5D方法。相比之下，3D方法直接利用由一系列二維切片或形式表示的完整的體積圖像。采用不同尺寸的直觀方法之后，采用多掃描、多路徑模型來捕獲圖像的下采樣特征。此外，為了充分利用三維環(huán)境和解決計算資源有限的問題，研究人員還研究了層次結構框架。
??有人提出在多個尺度或組合框架上提取特征。提出的多尺度框架來獲得胰腺分割中不同分辨率的信息減少了空間上下文和低分辨率條件下的問題，為三維醫(yī)學圖像的多層次分割提供了先導性研究。盡管這些網(wǎng)絡取得了成功，但其局限性在于它們在學習全局上下文和遠程空間依賴方面的性能較差，這可能會嚴重影響具有挑戰(zhàn)性任務的分割性能。

??Vision Transformers：視覺Transformers最近獲得了計算機視覺任務的關注。Dosovitskiy等人通過對純Transformers進行大規(guī)模的預訓練和微調(diào)，展示了圖像分類數(shù)據(jù)的最先進性能。在目標檢測中，端到端基于Transformers的模型在幾個基準測試中表現(xiàn)突出。最近提出了具有不同分辨率和空間嵌入的層次視覺Transformers。這些方法逐漸降低了Transformers層特征的分辨率，并利用了次采樣注意模塊。與這些方法不同的是，UNETR編碼器中表示的大小在所有Transformers層中保持固定，使用反卷積和卷積操作來改變提取特征的分辨率。
??最近，人們提出了多種方法，探索使用基于Transformers的模型進行二維圖像分割的可能性。Zheng等人引入了SETR模型，其中預訓練的Transformers編碼器與不同的基于CNN的解碼器被提出用于語義分割任務。Chenet al.提出了一種方法，通過在U-Net體系結構的瓶頸中使用Transformers作為額外的一層來實現(xiàn)多器官分割。Zhang等人建議將CNN網(wǎng)絡和Transformers分開使用，并融合它們的輸出。Valanarasu等人提出了一種基于Transformers的軸向注意機制用于二維醫(yī)學圖像分割。作者的模型與這些工作有關鍵的區(qū)別:
??(1) UNETR是為3D分割而量身定制的，并直接利用體積數(shù)據(jù);
??(2) UNETR使用Transformers作為分割網(wǎng)絡的主要編碼器，并通過跳過連接將其直接連接到解碼器，而不是使用它作為分割網(wǎng)絡中的注意層；
??(3) UNETR不依賴骨干CNN來生成輸入序列，而是直接利用標記化的patch。

??對于3D醫(yī)學圖像分割，Xie等人提出了一種框架，利用骨干CNN進行特征提取，利用Transformers處理編碼表示，利用CNN解碼器預測分割輸出。類似地，Wanget al.提出在3D編解碼器CNN的瓶頸處使用Transformers來完成腦腫瘤語義分割的任務。與這些方法不同的是，作者的方法通過使用跳躍連接直接將Transformers的編碼表示連接到解碼器。

模型方法:
??在圖2中展示了所提議的模型的概述。UNETR采用收縮-擴展模式，由一堆Transformer組成，編碼器通過跳躍連接到解碼器。Transformer工作在輸入嵌入的一維序列上，這是NLP中常用的方法。同樣，將三維輸入體積 ，分辨率(H,W,D)和C輸入通道劃分為平坦的均勻非重疊斑塊 ，創(chuàng)建一個一維序列，其中(P, P, P)表示每個patch的分辨率， 為序列的長度。



??圖2：UNETR架構概述。將一個三維輸入體(如MRI圖像的C= 4通道)劃分為一系列均勻的不重疊的patch，并使用線性層投影到嵌入空間中。該序列與位置嵌入一起添加，并用作Transformer模型的輸入。Transformer中不同層的編碼表示通過跳躍連接提取并與解碼器合并，以預測最終分割。

??圖1所示：UNETR的概述。提出的模型包括一個Transformers編碼器，它直接利用3D patch，并通過跳躍連接 連接到一個基于CNN的解碼器。

損失函數(shù)：
??損失函數(shù)是soft-dice損失和交叉熵損失的組合，它可以按體素方式計算：

其中I為體素的數(shù)量;J為類數(shù);Yi,j Gi,j分別表示j類在體素i上的概率輸出和one-hot編碼ground truth。

實驗與結果：
數(shù)據(jù)集：
??為了驗證方法的有效性，利用BTCV和MSD數(shù)據(jù)集在CT和MRI成像模式中進行三種不同的分割任務。

??BTCV (CT): BTCV數(shù)據(jù)集包括30名進行腹部CT掃描的受試者，其中13個器官由標注人員在范德比爾特大學醫(yī)學中心臨床放射科醫(yī)生的監(jiān)督下，進行標注。每次CT掃描均在門靜脈期進行造影增強，共80,225片 512×512 pixels，層厚1 ~ 6mm。通過將強度在[-1000,1000]HU到[0,1]的范圍內(nèi)歸一化，分別對每個體積進行預處理。在預處理過程中，將所有圖像重新采樣到各向同性體素間距1.0mm。多器官分割問題是一個13類1通道輸入的分割任務。

??MSD(MRI/CT):對于腦腫瘤的分割任務，使用整個訓練集484個多模態(tài)多位點MRI數(shù)據(jù)(FLAIR, T1w, T1gd, T2w)，并帶有膠質(zhì)瘤分割壞死/活性腫瘤和水腫的ground truth標簽進行模型訓練。本任務中MRI圖像體素間距為1.0 × 1.0 × 1.0 mm3。對體素強度進行z-score歸一化預處理。腦腫瘤分割問題是一個具有4通道輸入的3類分割任務。

??在脾臟分割任務中，使用了41個帶有脾臟體注釋的CT體積。任務9中卷的分辨率/間距范圍為0.613×0.613×1.50 mm3 ~ 0.977×0.977×8.0 mm3。在預處理過程中，將所有體積重新采樣到各向同性體素間距為1.0 mm。將圖像的體素強度按整體前景強度的第5和第95個百分位歸一化到[0,1]范圍。脾臟分割是一個具有1通道輸入的二進制分割任務。對于多器官和脾臟分割任務，我們對輸入的圖像隨機采樣，體積大小為[96,96,96]。對于腦分割任務，對輸入的圖像隨機采樣，體積大小為[128,128,128]。在所有的實驗中，前景/背景的隨機斑塊以1:1的比例采樣。

評價指標：
??使用Dice評分和95% Hausdorff Distance (HD)來評估分割的準確性。對于給定的語義類，讓Gi和 Pi表示體素i的地面真值和預測值，G’和P’分別表示ground真值和預測面點集。
Dice分數(shù)和HD指標被定義為:

95% HD使用地面真實值和預測表面點集之間距離的第95個百分位數(shù)。因此，在計算HD時，一個非常小的異常子集的影響是最小的。

實現(xiàn)細節(jié)：
??在pytorch和MONAI中實現(xiàn)UNETR。該模型使用NVIDIA DGX-1服務器進行訓練。所有模型都在批大小為6的情況下進行訓練，使用初始學習率為0.0001的AdamW優(yōu)化器進行20000次迭代。對于指定的批處理大小，20000次迭代的平均訓練時間為10小時?；赥ransformers的編碼器遵循ViT-B16架構，L= 12層，嵌入大小k = 768。使用的補丁分辨率為16×16×16。為了進行推斷，使用滑動窗口方法，在相鄰補丁之間的重疊部分為0.5，分辨率與第4.1節(jié)中指定的相同。沒有對Transformers主干使用任何預先訓練的權重(例如ImageNet上的ViT)，因為它沒有顯示出任何性能改進。對于BTCV數(shù)據(jù)集，已經(jīng)在其排行榜的標準和免費競賽中評估了模型和其他基線。自由競賽使用了來自同一隊列的額外數(shù)據(jù)，使訓練案例增加到80卷。對于所有的實驗，采用了5倍交叉驗證，比例為95:5。此外，使用了數(shù)據(jù)增強策略，如隨機旋轉(zhuǎn)90度、180度和270度，隨機翻轉(zhuǎn)軸位、矢狀位和冠狀位以及隨機尺度和偏移強度。使用集成來融合四個不同的五重交叉驗證模型的輸出。對于MSD數(shù)據(jù)集中的大腦和脾臟分割任務，將數(shù)據(jù)按照80:15:5的比例分成訓練、驗證和測試三部分。

定量評估：
??在BTCV排行榜上，UNETR在標準比賽和免費比賽中表現(xiàn)優(yōu)于最先進的方法。如表1所示，在自由競賽中，UNETR的整體平均Dice得分為0.899，比排名第二、第三和第四的方法分別高出1.238%、1.696%和5.269%。

??表1：BTCV測試集的分割性能的定量比較。頂部和底部分別代表標準比賽和自由比賽的基準。作者的方法與當前最先進的模型進行了比較。

??在標準競賽中，比較了UNETR與CNN和基于Transformers的基線的性能。UNETR實現(xiàn)了一個新的最先進的性能與平均骰子得分85.3%在所有器官。在脾、肝、胃等大器官上，方法在Dice得分上分別比次優(yōu)基線高出1.043%、0.830%和2.125%。此外，在小器官的分割中，在膽囊和腎上腺的Dice評分方面，方法顯著優(yōu)于第二最佳基線的6.382%和6.772%。
??在表2中，在MSD數(shù)據(jù)集上比較了UNETR與CNN以及基于Transformers的方法在腦腫瘤和脾臟分割任務中的性能。在腦分割方面，UNETR在所有語義類上的平均表現(xiàn)比最接近的基線高出1.5%。特別是，UNETR在分割腫瘤核心(TC)分區(qū)域方面表現(xiàn)得相當好。同樣，在脾臟分割方面，UNETR在Dice評分方面的表現(xiàn)比最佳競爭方法高出至少1.0%。

??表2：定量比較了MSD數(shù)據(jù)集在腦腫瘤和脾臟分割任務中的分割性能。WT、ET和TC分別表示全腫瘤、增強腫瘤和腫瘤核心子區(qū)域。

定性結果：
??多器官分割的定性比較如圖3所示。UNETR提高了腹部器官的分割性能。模型學習遠程依賴的能力在第3行(從上面開始)很明顯，其中nnUNet混淆了肝臟和胃組織，而UNETR成功地描繪了這些器官的邊界。在圖3中，第2行和第4行顯示腎臟和腎上腺對周圍組織的清晰檢測，這表明UNETR捕獲了更好的空間上下文。與基于二維Transformers的模型相比，UNETR展示了更高的邊界分割精度，因為它準確地識別了腎臟和脾臟之間的邊界?？梢娔懩以诘诙牛闻K和胃在第三排，門靜脈對肝臟在第五排。在圖4中，給出了在MSD數(shù)據(jù)集上對腦腫瘤分割的定性分割比較。具體來說，模型在捕捉腫瘤的細粒度細節(jié)方面表現(xiàn)得更好。

??圖3：BTCV交叉驗證中不同基線的定性比較。第一行為完整代表性CT切片。展示了4個放大的受試者(第2至5行)，方法顯示了腎臟和脾臟(第2行)、胰腺和腎上腺(第3行)、膽囊(第4行)和門靜脈(第5行)的分割效果。每個樣本上顯示了受試者的平均Dice得分。

??圖4：UNETR有效地捕捉細分輸出的細粒度細節(jié)。
??在所有數(shù)據(jù)集的實驗表明，UNETR的性能優(yōu)于CNN和基于Transformers的分割模型。具體來說，UNETR通過捕獲全局和局部依賴實現(xiàn)了更好的分割精度。在定性比較中，這一點在各種案例中得到了說明，在這些案例中，UNETR有效地捕捉了長期依賴關系(例如圖3中對胰腺尾部的精確分割)。
??此外，UNETR在BTCV排行榜上的細分表現(xiàn)顯示了新的最先進的基準，并驗證了其有效性。特別是對于小的解剖，UNETR優(yōu)于CNN和基于Transformers的模型。雖然3D模型已經(jīng)證明了對小器官(如膽囊、腎上腺)的高分割精度，但UNETR仍能以顯著優(yōu)勢領先于最佳競爭模型(見表1)。
這在圖3中也可以觀察到，其中UNETR對左腎上腺和右腎上腺的分割精度明顯更好，而且UNETR是唯一正確檢測腎上腺分支的模型。對于較具挑戰(zhàn)性的組織，如第4行膽囊和第5行門靜脈，它們與周圍肝組織的對比度較低，UNETR仍然能夠分割出清晰的連接邊界。

消融實驗：
??解碼器的選擇：在表3中，通過比較UNETR和其他解碼器架構在兩個MRI和CT模式的代表性分割任務上的性能來評估解碼器的有效性。在這些實驗中，使用UNETR的編碼器，但用樸素向上采樣(NUP)、漸進向上采樣(PUP)和多尺度聚合(MLA)的3D對應版本替換解碼器。觀察到這些解碼器架構產(chǎn)生次優(yōu)性能，盡管MLA略微優(yōu)于NUP和PUP。對于腦腫瘤的分割，UNETR比MLA、PUP和NUP解碼器的平均準確率分別高出2.7%、4.3%和7.5%。同樣，對于脾臟分割，UNETR外形成的MLA、PUP和NUP分別為1.4%、2.3%和3.2%。

??表3：解碼器架構對分割性能的影響。（Naive UpSampling、Progressive UpSampling and Multi-scale Aggregation）

??Patch Resolution：較低的輸入補丁分辨率導致較大的序列長度，因此內(nèi)存消耗更高，因為它與分辨率的立方成反比。如表4所示，實驗表明，降低分辨率可以持續(xù)提高性能。具體來說，將patch分辨率從32降低到16，在脾臟和大腦分割任務的平均Dice得分方面分別提高了1.1%和0.8%的性能。由于內(nèi)存的限制，沒有進行低分辨率的實驗。

??表4：patch分辨率對分割性能的影響。

??模型與計算復雜度：在表5中，給出了BTCV基準中模型的FLOPs次數(shù)、參數(shù)和平均推理時間。浮點數(shù)和推斷時間是基于96×96×96的輸入大小并使用滑動窗口方法計算的。根據(jù)基準，UNETR是一個中等大小的模型，具有92.58M參數(shù)和41.19G浮點數(shù)。相比而言，其他基于Transformers的方法，如CoTr、TransUNet和SETR，分別有46.51M、96.07M和86.03M參數(shù)和399.21G、48.24G和43.49G FLOPs。UNETR顯示了類似的模型復雜性，同時在BTCV基準中大大優(yōu)于這些模型。nnUNet和ASPP基于cnn的分割模型分別有19.07M和47.92M參數(shù)，412.65G和44.87G FLOPs。同樣，UNETR優(yōu)于這些基于cnn的模型，同時具有中等的模型復雜性。此外，UNETR的平均推理時間僅次于nnUNet，而且明顯快于基于Transformers的模型，如SETR， TransUNet和CoTr。

??表5所示：BTCV實驗中各模型的參數(shù)數(shù)、次數(shù)和平均推斷時間的比較。

結論：
??本文介紹了一種新穎的基于Transformer的結構，稱為UNETR，用于體積醫(yī)學圖像的語義分割，將該任務重新定義為一個一維序列到序列的預測問題。作者建議使用Transformers編碼器來提高模型的能力，以學習遠程依賴關系，并在多個尺度上有效地捕獲全局上下文表示。
??驗證了UNETR在CT和MRI不同體積分割任務中的有效性。在BTCV多器官分割排行榜上，UNETR競賽中取得了最新的水平表現(xiàn)，并在MSD數(shù)據(jù)集上優(yōu)于腦腫瘤和脾臟分割的競爭方法。該方法可作為醫(yī)學圖像分析中一類新的基于Transformers的分割模型的基礎。文章來源地址http://www.zghlxwxcb.cn/news/detail-821708.html

到了這里，關于[深度學習論文筆記]UNETR: Transformers for 3D Medical Image Segmentation的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關文章，希望大家以后多多支持TOY模板網(wǎng)！