超越nnFormer！UNETR++：高效準確的3D醫學影像分割

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達

點選進入—>醫療影像微信技術交流群

轉載自：CVHub

UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation

論文連結：

https://arxiv.org/abs/2212.04497

程式碼連結：

https://github.com/Amshaker/unetr_plus_plus

導讀

這篇論文主要講述了一種名為 UNETR++ 的 3D 醫學影像分割方法，它提供了高質量的分割結果，並具有高效的引數和計算成本。作者介紹了一種新的有效的配對注意力（EPA）模組，該模組使用一對基於空間和通道注意的相互依賴分支來有效地學習空間和通道方向的區分性特徵。實驗結果表明，該方法在 Synapse、BTCV 和 ACDC 資料集上均優於現有方法。

背景

3D 分割是醫學影像中的一個基礎問題，並且用於許多應用，包括腫瘤識別和器官定位等診斷目的。3D 分割任務通常採用 U-Net 類似的編碼器-解碼器架構，其中編碼器生成 3D 影像的分層低維表示，解碼器將這個學習的表示對映到體素分割。早期的基於 CNN 的方法在編碼器和解碼器中分別使用卷積和反捲積，但很難實現準確的分割結果，可能是由於其侷限的區域性感受野。另一方面，基於 transformer 的方法天生是全域性的，並且最近已經證明了具有競爭效能的成本，但模型複雜度增加。這篇論文提出了一種混合架構來結合區域性卷積和全域性注意力的優點。作者指出，這種架構能夠在提高分割精度的同時減小模型的引數數量和 FLOPs，提高模型的魯棒性。作者認為，這種方法的優勢在於它捕捉了空間和通道特徵之間的顯式依賴關係，從而提高了分割質量。

UNETR++是一種新型的3D醫學影像分割混合層次架構，旨在提高分割準確度和效率。它是在UNET的基礎上提出的一種改進模型，透過引入精細化模組和跨層連線來增強特徵表示能力，並透過尺度匹配和跨層金字塔來減少模型引數和計算複雜度。經過實驗驗證，UNETR++在主流資料集上的分割準確度和效率都優於其他現有模型。UNETR++基於最近提出的UNETR框架，引入了一種新的有效的配對注意力（EPA）模組，透過在兩個分支中應用空間注意力和通道注意力，有效地捕捉了豐富的相互依存的空間和通道特徵。我們的EPA中的空間注意力將key和value投影到固定的較低維空間，self attention的計算與input token的數量成線性關係。這有助於提高模型的特徵表示能力，同時又不會增加太多的引數和計算複雜度。另一方面，我們的通道注意力透過在通道維度上對query和key執行點積運算，強調了通道特徵對映之間的依賴性。此外，為了捕捉空間和通道特徵之間的強相關性，query和key的權重在兩個分支之間共享，這也有助於控制網路引數的數量。相比之下，值的權重是獨立的，以強制在兩個分支中學習互補特徵。這種設計能夠有效地平衡兩個分支的貢獻，提高模型的效能。

Method

上圖展示了UNETR++模型的層級編解碼結構。

動機： 在設計混合框架時要考慮的兩個理想屬性：效率和準確性。這兩個屬性都非常重要，因為它們可以幫助框架在實際應用中更好地發揮作用。在設計混合框架時，應該儘量將這兩個屬性結合起來，以達到更好的效能。
有效的全域性注意力： 現有混合方法中使用的自注意力操作複雜度為平方，並且在體積醫學分割中具有較高的計算代價。與這些方法不同，作者認為，在特徵通道而不是體積維度上計算自注意力可以將複雜度從平方降低到線性。此外，空間注意力資訊可以透過將key和query空間矩陣投影到較低維空間中來有效地學習。這種方法可以在保證效能的同時減少計算複雜度，從而更好地處理大型體積資料。
增強空間和通道特徵表示： 大多數現有的體積醫學影像分割方法通常透過注意力計算來捕獲空間特徵，並忽略了通道資訊。有效地結合空間維度中的互動作用和通道特徵之間的相互依賴關係，有望提供豐富的上下文空間通道特徵表示，從而提高掩模預測的準確性。透過這種方式，演算法可以更好地利用通道資訊來捕獲更豐富的空間資訊，從而提高分割精度。

本文提出了一種UNETR++框架，旨在更有效地學習空間通道特徵表示。為了實現這一目標，這篇論文中提出了一種新的EPA模組，它包含兩個注意力模組，透過共享key-query方案來有效地在空間和通道維度上編碼資訊。該論文還提出了在編碼階段和解碼階段之間透過跳過連線來連線的方法，以便在不同的解析度上合併輸出。這有助於恢復在下采樣操作期間丟失的空間資訊，從而預測更精確的輸出。與編碼器類似，解碼器由四個階段組成，每個階段都包含一個使用反捲積來增加特徵圖解析度的上取樣層，然後是 EPA 模組（除了最後一個解碼器）。在每兩個解碼器階段之間，通道數減少了一倍。因此，最後一個解碼器的輸出與卷積特徵圖融合，以恢復空間資訊並增強特徵表示。然後將得到的輸出輸入 3×3×3 和 1×1×1 卷積塊，生成體素方面的最終掩碼預測。

接下來是 EPA 模組的介紹。EPA 模組執行高效的全域性注意力，有效地捕獲豐富的空間-通道特徵表示。EPA 模組包括空間注意力和通道注意力模組。空間注意力模組將自我關注的複雜度從二次降低到線性。另一方面，通道注意力模組有效地學習通道特徵圖之間的相互依賴性。EPA 模組基於兩個注意力模組之間的共享key和query方案，以便互相資訊交流，以生成更好、更高效的特徵表示。這可能是因為透過共享key和query來學習互補特徵，但使用不同的值層。

實驗

在abdominal multi-organ Synapse dataset中，對比SOTA模型，UNETR++僅用了其1/3的引數量和1/4的計算量取得了最高的精度。

在多器官分割的BTCV測試集中，UNETR++的綜合精度比SOTA模型高出1.5%以上。

在ACDC資料集中，UNETR++的綜合精度比SOTA模型高出0.2%以上。

將EPA模組分別應用在編碼器和解碼器的精度提升。

效果展示

可以很清晰地觀察到，無論是小體積還是粘連等複雜場景，UNETR++都能略勝一籌。

侷限性

本文所提出的UNETR++在三種醫學影像分割資料集上均以更少的模型複雜度獲得了更高的分割精度，超越了目前所有的醫學影像分割SOTA模型。然而，在對於以下兩種情況：器官形狀模糊以及器官邊界不明顯，UNETR++也無法很好的分割。這是由於這些困難場景樣本遠少於正常場景樣本，導致模型無法很好地關注到困難樣本特徵。未來，作者將會在資料預處理階段使用特定的幾何資料增強技術來對這些困難場景進行針對性的最佳化。

總結

本文提出了一種用於3D醫學圖象分割的分層方法：UNETR++。透過引入有效的配對注意力（EPA）模組，UNETR++能夠豐富空間和通道之間相關依賴的特徵。
在配對注意力（EPA）模組中，共享query和key的對映函式權重能夠使得空間和通道分支進行更高效地交流特徵資訊，這可以為雙方提供互補的有益特徵並且降低引數量。
UNETR++在三個國際著名醫學圖象分割資料集中均以更少的模型複雜度取得最更高的精度結果。

點選進入—>CV微信技術交流群

CVPR/ECCV 2022論文和程式碼下載

後臺回覆：CVPR2022，即可下載CVPR 2022論文和程式碼開源的論文合集

後臺回覆：ECCV2022，即可下載ECCV 2022論文和程式碼開源的論文合集

後臺回覆：Transformer綜述，即可下載最新的3篇Transformer綜述PDF

醫療影像交流群成立

掃描下方二維碼，或者新增微信：CVer222，即可新增CVer小助手微信，便可申請加入CVer-醫療影像微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer等。

一定要備註：研究方向+地點+學校/公司+暱稱（如醫療影像+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群