NVIDIA SWI UNETR模型在醫(yī)療中的應(yīng)用

NVIDIA 數(shù)據(jù)科學(xué)家本周在享有盛譽(yù)的MICCAI 2021醫(yī)學(xué)成像會(huì)議上，在腦腫瘤分割挑戰(zhàn)驗(yàn)證階段占據(jù)了前 10 名的三位。

現(xiàn)在已經(jīng)進(jìn)入第十個(gè)年頭， BraTS 挑戰(zhàn)項(xiàng)目要求申請(qǐng)者提交最先進(jìn)的人工智能模型，用于在多參數(shù)磁共振成像（ mpMRI ）研究中分割異質(zhì)性腦膠質(zhì)母細(xì)胞瘤亞區(qū)域，這是一項(xiàng)極具挑戰(zhàn)性的任務(wù)。

參與者還可以關(guān)注分類方法的第二項(xiàng)任務(wù)，以預(yù)測(cè) MGMT 啟動(dòng)子甲基化狀態(tài)。

2000 多個(gè) AI 模型被提交給了挑戰(zhàn)，該挑戰(zhàn)由醫(yī)學(xué)圖像計(jì)算和計(jì)算機(jī)輔助干預(yù)學(xué)會(huì)、北美放射學(xué)會(huì)和美國(guó)神經(jīng)放射學(xué)會(huì)聯(lián)合組織。

NVIDIA 開發(fā)者在挑戰(zhàn)驗(yàn)證階段排名第一、第二和第七，每個(gè)人都創(chuàng)建了不同類型的用于腫瘤分割的 AI 模型方法，包括優(yōu)化的 U-Net 模型、具有自動(dòng)超參數(shù)優(yōu)化的 SegResNet 模型和基于 transformer 的計(jì)算機(jī)視覺方法的 Swin-UNETR 模型。

NVIDIA 獲獎(jiǎng)?wù)叨际褂昧碎_源 PyTorch 框架MONAI（人工智能醫(yī)療開放網(wǎng)絡(luò)），這是一個(gè)由學(xué)術(shù)界和行業(yè)領(lǐng)袖構(gòu)建的免費(fèi)、社區(qū)支持的計(jì)劃，旨在將醫(yī)療成像深度學(xué)習(xí)的最佳實(shí)踐標(biāo)準(zhǔn)化。

用于腦腫瘤分割的優(yōu)化 U-Net –排名# 1

該優(yōu)化 DU 網(wǎng)絡(luò)模型是一種編碼器 – 解碼器類型的卷積網(wǎng)絡(luò)體系結(jié)構(gòu)，用于快速、精確的圖像分割，它首先處于 BraTS 驗(yàn)證階段。它的標(biāo)準(zhǔn)化統(tǒng)計(jì)排名得分為 0 。 267 。

設(shè)計(jì)優(yōu)化 U 型網(wǎng)絡(luò)的起點(diǎn)是 BraTS 2020 獲獎(jiǎng)解決方案：神經(jīng)網(wǎng)絡(luò)在腦腫瘤分割中的應(yīng)用。該團(tuán)隊(duì)的目標(biāo)是優(yōu)化 U-Net 體系結(jié)構(gòu)以及培訓(xùn)計(jì)劃。為了找到最佳的模型結(jié)構(gòu)，數(shù)據(jù)科學(xué)家進(jìn)行了廣泛的燒蝕研究，發(fā)現(xiàn)香草 U-Net 和深度監(jiān)督產(chǎn)生了最好的結(jié)果。

進(jìn)一步優(yōu)化 U-Net 模型，在輸入端添加額外通道，對(duì)前景體素進(jìn)行一次熱編碼，將編碼器深度和卷積通道數(shù)量增加一個(gè)級(jí)別。通過在較低的解碼器級(jí)別上添加兩個(gè)額外的輸出頭，可以更好地實(shí)現(xiàn)梯度流和更準(zhǔn)確的預(yù)測(cè)，從而在深度監(jiān)督下更好更快地訓(xùn)練模型。MONAI用于數(shù)據(jù)預(yù)處理，以清除數(shù)據(jù)并去除數(shù)據(jù)中的噪聲，以及推斷模型。 NVIDIA 數(shù)據(jù)加載庫（ DALI ）使用并支持?jǐn)?shù)據(jù)擴(kuò)充，這是一種用于人為擴(kuò)大數(shù)據(jù)集大小的技術(shù)，它通過將數(shù)據(jù)擴(kuò)充卸載到 GPU 來解決 CPU 瓶頸問題。

帶有最新 PyTorch、cuDNN 和 CUDA 版本的 NVIDIA PyTorch 容器用于優(yōu)化 U-Net 模型以進(jìn)行快速訓(xùn)練。自動(dòng)混合精度 （AMP） 用于將 AI 模型的內(nèi)存占用減少 2 倍并加快訓(xùn)練速度。訓(xùn)練在 8 個(gè) NVIDIA A100 GPU 上完成 1000 個(gè) epoch，與 V100 GPU 相比，速度提高了 2 倍。幾乎 100% 的 GPU 被使用，展示了網(wǎng)絡(luò)優(yōu)化以高效使用 GPU。此 3D U-Net 模型可用于任何 3D 模式，例如 MRI 和 CT。了解更多 GPU 高效的 nnU-Net 實(shí)現(xiàn)。

圖 1 。挑戰(zhàn)驗(yàn)證數(shù)據(jù)集上的預(yù)測(cè)。在第一行， T2 模態(tài)是可視化的。第二行顯示以下顏色的模型預(yù)測(cè)：紫色 – 背景、藍(lán)色 -NCR 、綠松石色 -ED 、黃色 -ET 。

SegResNet ： 3D 大腦磁共振成像語義分割中的冗余減少–排名# 2

該方法在 BRAT 挑戰(zhàn)中排名第二，將在 MONAI 中提供，基于 MONAI 組件，旨在展示其應(yīng)用的實(shí)用性和靈活性。主要模型是 MONAI 中的 SegResNet 架構(gòu)，這是一種基于標(biāo)準(zhǔn)編碼器 – 解碼器的卷積神經(jīng)網(wǎng)絡(luò)（ CNN ），類似于 U-Net 。該方法是 MONAI 自動(dòng)化（ AutoML ）計(jì)劃的一部分，使用超參數(shù)優(yōu)化和調(diào)優(yōu)自動(dòng)選擇超參數(shù)。

該方法在驗(yàn)證階段的排行榜（ team NVAUTO ）中取得了最佳性能，并在基于個(gè)案排名和擾動(dòng)分析的綜合排名中取得了 0 。 272 的排名。組織者表示，該排名與第一排名的解決方案在統(tǒng)計(jì)上沒有顯著差異，兩種方法被認(rèn)為在統(tǒng)計(jì)上相似。

這項(xiàng)工作增加了兩項(xiàng)新貢獻(xiàn)，以進(jìn)一步提高性能。首先，對(duì)訓(xùn)練過程進(jìn)行修改，以在學(xué)習(xí)的特征表示上強(qiáng)制執(zhí)行某些屬性。通過借鑒自監(jiān)督文獻(xiàn)中的思想，將特征維度正則化，使不同解剖區(qū)域之間的冗余度最小。同時(shí)，相同解剖結(jié)構(gòu)的區(qū)域被鼓勵(lì)相似。這允許更好的網(wǎng)絡(luò)行為和泛化。其次，采用自適應(yīng)融合技術(shù)自適應(yīng)地選擇模型子集進(jìn)行融合。這有助于避免某些模型預(yù)測(cè)中的潛在異常值，并進(jìn)一步提高最終的集合性能。

該方法在 MONAI 中實(shí)現(xiàn)，以 PyTorch 為基礎(chǔ)，并在四個(gè) NVIDIA V100 GPU 上使用骰子損失函數(shù)進(jìn)行了為期 16 小時(shí)的 300 個(gè)歷元的訓(xùn)練。該團(tuán)隊(duì)使用 NVIDIA 提供的 PyTorch 容器和 AMP 在超參數(shù)優(yōu)化過程中實(shí)現(xiàn)快速訓(xùn)練。該方法通過 5 倍交叉驗(yàn)證進(jìn)行訓(xùn)練，并從每個(gè)折疊中保留表現(xiàn)最佳的檢查點(diǎn)。總的來說，保存了 25 個(gè)模型檢查點(diǎn)，但使用自適應(yīng)加密，只有一半用于最終預(yù)測(cè)。由于該方法是完全基于 CNN 的，推理時(shí)間很快，并且可以在整個(gè)輸入圖像上一步完成，無需任何滑動(dòng)窗口。單模型推理只需不到一秒鐘的時(shí)間，這使得高通量和幾乎實(shí)時(shí)的結(jié)果成為可能，這在臨床環(huán)境中非常重要。

圖 2 。一個(gè)典型的分割示例，預(yù)測(cè)的標(biāo)簽覆蓋在 MRI 軸向、矢狀面和冠狀面切片上。整個(gè)腫瘤（ WT ）類別包括所有

可見標(biāo)簽（綠色、黃色和紅色標(biāo)簽的結(jié)合），腫瘤核心（ TC ）類別為

紅色和黃色的結(jié)合，增強(qiáng)腫瘤核心（ ET ）級(jí)別顯示為黃色

（過度活躍的腫瘤部分）。

Swin UNETR ：用于腦腫瘤三維語義分割的移位窗口 transformers –排名第 7

在 BraTS 挑戰(zhàn)賽中排名第七的 Swin UNETR 是一個(gè)基于 transformer 的模型，而不是 CNN 模型。在 MONAI 中實(shí)現(xiàn)，它在整個(gè)腫瘤、腫瘤核心和增強(qiáng)腫瘤分割類中的平均 Dice 得分為 92 。 94% ， Ha USD orff 距離為 1 。 7 。

transformers 是一類新的基于深度學(xué)習(xí)的模型，用于序列到序列的預(yù)測(cè)任務(wù)。在其原始公式中，它們由編碼器和解碼器組成。編碼器由多個(gè)層次組成，多個(gè)層次的感知器跟隨多個(gè)層次的自我注意。每個(gè)模塊的輸入通過剩余塊添加到輸出，并進(jìn)行歸一化。自我注意層學(xué)習(xí)從隱藏層計(jì)算的值的加權(quán)和，可以突出給定輸入序列的重要特征。雖然它們最初被提議用于自然語言處理中的機(jī)器翻譯任務(wù)，但后來它們已成功應(yīng)用于計(jì)算機(jī)視覺和蛋白質(zhì)藥物生成等其他領(lǐng)域，并在各種基準(zhǔn)測(cè)試中取得了最先進(jìn)的性能。在計(jì)算機(jī)視覺中， transformer s 在各種基準(zhǔn)方面取得了新的最先進(jìn)的表現(xiàn)。 Swin UNETR 是一種利用 GPU 體系結(jié)構(gòu)和性能不斷進(jìn)步來構(gòu)建模型的新方法。

Swin transformers 是分層 transformer s ，其表示是使用移位窗口（ Swin ）計(jì)算的。這些 transformer 非常適合計(jì)算機(jī)視覺任務(wù)，如目標(biāo)檢測(cè)、圖像分類、語義分割等。 Swin transformers 可以更有效地模擬兩個(gè)域之間的差異，例如對(duì)象比例的變化和圖像中像素的高分辨率，并且可以作為通用的視覺管道。該 NVIDIA SWN UNETR 模型利用 SWN transformer 編碼器直接利用輸入數(shù)據(jù)的 3D 補(bǔ)丁，而不依賴 CNN 進(jìn)行特征提取。這使得 Swin UNETR 能夠訪問輸入數(shù)據(jù)中的上下文多模態(tài)信息，并將其作為標(biāo)記化嵌入 transformer 編碼器進(jìn)行有效處理。然后，基于 transformer 的 Swin UNETR 編碼器通過跳過連接連接到 U 形架構(gòu)中的 CNN 解碼器，以進(jìn)行最終分段預(yù)測(cè)。

圖 3 。使用 BraTS2021 驗(yàn)證數(shù)據(jù)集，覆蓋在 4 種單獨(dú)成像模式上的 TC 、 WT 和 ET 區(qū)域的 Swin UNETR 分割輸出。

NVIDIA 的 SWI UNETR 模型在 NVIDIA DGX-1 團(tuán)簇上使用八個(gè) GPU 進(jìn)行訓(xùn)練，初始學(xué)習(xí)率為 0 。 0008 ，并使用 AdamW 優(yōu)化算法。使用 128 * 128 * 128 輸入數(shù)據(jù)的隨機(jī)面片以及隨機(jī)軸鏡像翻轉(zhuǎn)和強(qiáng)度偏移的數(shù)據(jù)增強(qiáng)策略。 每輪培訓(xùn)需要 24 小時(shí)才能完成。與常用的基于 CNN 的分割模型相比， Swin UNETR 在失敗次數(shù)方面更有效，在可訓(xùn)練參數(shù)數(shù)量方面具有中等的模型復(fù)雜度。它可以被有效地訓(xùn)練和用于推理。對(duì)于模型優(yōu)化，使用一個(gè)通用的軟骰子損失函數(shù)學(xué)習(xí)分割不同的腦腫瘤區(qū)域，每個(gè)類有一個(gè)單獨(dú)的輸出通道。在整個(gè) BRATS21 訓(xùn)練集上使用 5 倍交叉驗(yàn)證方案對(duì)模型進(jìn)行訓(xùn)練，其中通過平均 2 個(gè)不同 5 倍交叉驗(yàn)證的 10 個(gè)模型的輸出來計(jì)算最終分割輸出。

從 MONAI 開始。

在這里查找 BRAT 參與者的完整排行榜。

關(guān)于作者

Vanessa Braunstein 在 NVIDIA 的醫(yī)療團(tuán)隊(duì)從事產(chǎn)品營(yíng)銷工作。此前，她在基因組學(xué)、醫(yī)學(xué)成像、制藥、化學(xué)和診斷公司從事產(chǎn)品開發(fā)和營(yíng)銷。她學(xué)習(xí)分子和細(xì)胞生物學(xué)、公共衛(wèi)生和商業(yè)。

審核編輯：郭婷