新型transformer模型實現(xiàn)3D醫(yī)學(xué)圖像分析最新基準(zhǔn)

在計算機視覺和模式識別會議（ CVPR ）上， NVIDIA 研究人員發(fā)表了 35 多篇論文。這包括對 移動窗口 UNEt TRansformers （ Swin-UNETR ）的研究，這是第一個基于變壓器的預(yù)訓(xùn)練框架，專為 3D 醫(yī)學(xué)圖像分析中的自我監(jiān)督任務(wù)而設(shè)計。這項研究是創(chuàng)建用于數(shù)據(jù)注釋的預(yù)訓(xùn)練、大規(guī)模和自監(jiān)督三維模型的第一步。

作為一種基于 transformer 的計算機視覺方法， Swin UNETR 采用了 MONAI ，這是一種開源的 PyTorch 框架，用于深入學(xué)習(xí)醫(yī)療成像，包括放射學(xué)和病理學(xué)。使用這種預(yù)訓(xùn)練方案， Swin UNETR 為各種醫(yī)學(xué)圖像分割任務(wù)設(shè)定了新的最先進(jìn)的基準(zhǔn)，并一致證明了其有效性，即使只有少量的標(biāo)記數(shù)據(jù)。

Swin UNETR 模型培訓(xùn)

Swin UNETR 模型在 NVIDIA DGX-1 集群 使用八個 GPU 和 AdamW 優(yōu)化算法。對 5050 張來自健康和不健康受試者不同身體部位的公開 CT 圖像進(jìn)行預(yù)訓(xùn)練，以保持?jǐn)?shù)據(jù)集的平衡。

對于 3D Swin transformer 編碼器的自我監(jiān)督預(yù)訓(xùn)練，研究人員使用了各種借口任務(wù)。隨機裁剪的標(biāo)記使用不同的變換（如旋轉(zhuǎn)和剪切）進(jìn)行增強。這些標(biāo)記用于掩蔽體修復(fù)、旋轉(zhuǎn)和對比學(xué)習(xí)，用于編碼器學(xué)習(xí)訓(xùn)練數(shù)據(jù)的上下文表示，而不會增加數(shù)據(jù)注釋的負(fù)擔(dān)。

圖 1 ：。預(yù)培訓(xùn)框架概述。輸入的 CT 圖像被隨機裁剪成子體積，并通過隨機內(nèi)部剪切和旋轉(zhuǎn)進(jìn)行增強，然后作為輸入饋送到 Swin-UNETR 編碼器

Swin UNETR 背后的技術(shù)

Swin Transformers 采用分層視覺 transformer （ ViT ）進(jìn)行非重疊窗口的局部自我注意計算。這打開了為大型公司創(chuàng)建醫(yī)療專用 ImageNet 的機會，消除了創(chuàng)建醫(yī)療 AI 模型需要大量高質(zhì)量注釋數(shù)據(jù)集的瓶頸。

與 CNN 體系結(jié)構(gòu)相比， ViT 在從未標(biāo)記數(shù)據(jù)（數(shù)據(jù)集越大，預(yù)訓(xùn)練主干越強）進(jìn)行全局和局部表示的自監(jiān)督學(xué)習(xí)方面表現(xiàn)出非凡的能力。用戶可以在下游任務(wù)（例如，分割、分類和檢測）中使用極少量的標(biāo)記數(shù)據(jù)微調(diào)預(yù)訓(xùn)練模型。

這種體系結(jié)構(gòu)在本地窗口中計算自我注意，與 ViT 相比表現(xiàn)出更好的性能。此外， Swin Transformers 的層次性使其非常適合需要多尺度建模的任務(wù)。

繼領(lǐng)先的 UNETR 模型成功使用直接使用 3D 補丁嵌入的基于 ViT 的編碼器之后， Swin UNETR 使用了具有金字塔結(jié)構(gòu)的 3D Swin transformer 編碼器。

在 Swin UNETR 的編碼器中，由于計算簡單的全局自我注意對于高分辨率特征地圖是不可行的，因此在本地窗口中計算自我注意。為了增加局部窗口以外的感受野，使用窗口移位來計算不同窗口的區(qū)域相互作用。

Swin UNETR 的編碼器通過跳過連接以五種不同的分辨率連接到剩余的類似 UNet 的解碼器。它可以為密集的預(yù)測任務(wù)（如醫(yī)學(xué)圖像分割）捕獲多尺度特征表示。

Swin UNETR 模型性能

在對 CT 中的 13 個腹部器官和 醫(yī)學(xué)分段十項全能（ MSD ） 數(shù)據(jù)集中的分割任務(wù)使用 超越顱穹窿（ BTCV ）分割挑戰(zhàn) 進(jìn)行微調(diào)后，該模型在公共排行榜上達(dá)到了最先進(jìn)的精度。

BTCV

在 BTCV 中， SwinUnetr 的平均骰子數(shù)為 0.918 ，優(yōu)于其他排名靠前的模型。

圖 2 ：。 BTCV 挑戰(zhàn)多器官分割排行榜骰子結(jié)果

對于較小的器官，如脾靜脈和門靜脈（ 3.6% ）、胰腺（ 1.6% ）和腎上腺（ 3.8% ），與之前的最先進(jìn)方法相比，有了改進(jìn)小器官數(shù)據(jù)標(biāo)簽分割對于放射科醫(yī)生來說是一項極其困難的任務(wù)。

通訊簿標(biāo)簽

在 MSD 中， Swin UNETR 在腦腫瘤、肺、胰腺和結(jié)腸方面取得了最先進(jìn)的表現(xiàn)。心臟、肝臟、海馬、前列腺、肝血管和脾臟的結(jié)果具有可比性?？偟膩碚f， Swin UNETR 在所有 10 項任務(wù)中的平均骰子率為 78.68% ，是最好的，并在 MSD 排行榜上排名第一。

圖 4 ：。使用 MSD 數(shù)據(jù)集進(jìn)行肝血管分割任務(wù)的定性輸出。你可以看到， Swin UNETR 比 SOTA DiNTS 跑贏了 3% 。這一變化很重要，因為它會導(dǎo)致更準(zhǔn)確的分段輸出

與DiNTS（一種用于醫(yī)學(xué)圖像分割的強大AutoML方法）相比，Swin UNETR使用更少的訓(xùn)練時間，顯示出更好的分割性能。例如，肝血管分割任務(wù)的定性分割輸出證明了Swin UNETR能夠更好地建模長期空間依賴性。

圖 5 ：。 Swin UNETR 模型體系結(jié)構(gòu)由一個 Swin transformer 編碼器組成，該編碼器使用 3D 補丁，并通過不同分辨率的跳過連接連接到基于 CNN 的解碼器

結(jié)論

Swin UNETR 體系結(jié)構(gòu)在使用變壓器的醫(yī)療成像方面提供了急需的突破。鑒于醫(yī)學(xué)成像需要快速構(gòu)建準(zhǔn)確的模型， Swin UNETR 體系結(jié)構(gòu)使數(shù)據(jù)科學(xué)家能夠?qū)Υ罅课礃?biāo)記數(shù)據(jù)進(jìn)行預(yù)訓(xùn)練。這減少了放射科醫(yī)生、病理學(xué)家和其他臨床團隊進(jìn)行專家注釋的成本和時間。這里我們展示了用于器官檢測和自動體積測量的 SOTA 分割性能。

審核編輯：郭婷