Title: 3D UX-Net: a Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation
Author: Ho Hin Lee et al. (范德堡大學(xué))
Paper: https://arxiv.org/abs/2209.15076
Github: https://github.com/MASILab/3DUX-Net
引言
眾所周知,大多數(shù)醫(yī)學(xué)圖像如 MRI 和 CT 是屬于 volumetric data 類型。因此,為了更加充分的利用體素信息,近幾年已經(jīng)提出了不少 3D CNNs 的模型,如 SwinUNETR
、UNETR
以及筆者前段時(shí)間分享過的UNETR++
等。
整體來說,這些模型性能是越來越高,在幾個(gè)主流的 3D 數(shù)據(jù)基準(zhǔn)測試中也實(shí)現(xiàn)了大大小小的 SOTA,特別是 3D 醫(yī)學(xué)圖像分割這塊。當(dāng)然,時(shí)代在進(jìn)步,作為一名高科技前沿從業(yè)者本身也是需要不斷汲取新的知識(shí)營養(yǎng)才能不被輕易的淘汰。今天小編就帶大家解讀下 ICLR 2023 新鮮出爐的 3D 醫(yī)學(xué)圖像分割之星——3D UX-Net
。
3D UXNet
是一種輕量級3D卷積神經(jīng)網(wǎng)絡(luò),其使用 ConvNet
模塊調(diào)整分層 Transformer
以實(shí)現(xiàn)穩(wěn)健的體素分割,在三個(gè)具有挑戰(zhàn)性的腦體積和腹部成像公共數(shù)據(jù)集與當(dāng)前的 SOTA 模型如 SwinUNETR
對比,同時(shí)在以下三大主流數(shù)據(jù)集性能均達(dá)到了最優(yōu):
-
MICCAI Challenge 2021 FLARE
-
MICCAI Challenge 2021 FeTA
-
MICCAI Challenge 2022 AMOS
具體的,相比于 SwinUNETR,3D UX-Net 將 Dice 從 0.929 提高到 0.938 (FLARE2021),Dice 從 0.867 提高到 0.874 (Feta2021)。此外,為了進(jìn)一步評估 3D UX-Net 的遷移學(xué)習(xí)能力,作者在 AMOS2022
數(shù)據(jù)集上訓(xùn)練得到的模型在另一個(gè)數(shù)據(jù)集上取得了 2.27% Dice 的提升(0.880 → 0.900)。
動(dòng)機(jī)
先來看看最近提出的一眾基于 Transformer
架構(gòu)的 ViT
模型有什么優(yōu)劣勢。以 SwinUNETR
為例,其將 ConvNet
的一些先驗(yàn)引入到了 Swin Transformer
分層架構(gòu)中,進(jìn)一步增強(qiáng)了在 3D 醫(yī)學(xué)數(shù)據(jù)集中調(diào)整體素分割的實(shí)際可行性。此類 "Conv+Transforemr" 的組合拳的有效性在很大程度上歸功于以下兩個(gè)因素:
- 非局部自注意力所帶來的大感受野
- 大量的模型參數(shù)
為此,本文作者想到了應(yīng)用深度卷積以更少的模型參數(shù)來模擬這些方法讓網(wǎng)絡(luò)學(xué)會(huì)如何捕獲更大感受野的行為。這一點(diǎn)筆者在前文《關(guān)于語義分割的億點(diǎn)思考》中也提過,文中對整個(gè)任務(wù)進(jìn)行了深度剖析,有興趣的讀者可以自行翻閱公眾號(hào)歷史文章:
語義分割任務(wù)的核心思想是如何高效建模上下文信息,它是提升語義分割性能最為重要的因素之一,而有效感受野則大致決定了網(wǎng)絡(luò)能夠利用到多少上下文信息。
回到正文,本文的設(shè)計(jì)思路主要是受 ConvNeXt
啟發(fā),還沒了解過的同學(xué)也可自行移動(dòng)到公眾號(hào)翻閱歷史文章,筆者之前對該系列也講解過了,此處不再詳述。3D UX-Net
核心理念是設(shè)計(jì)出一種簡單、高效和輕量化的 網(wǎng)絡(luò),其適用于 hierarchical transformers
的能力同時(shí)保留使用 ConvNet
模塊的優(yōu)勢,如歸納偏置。具體地,其編碼器模塊設(shè)計(jì)的基本思想可分為: 1) block-wise(微觀層面) 及 2) layer-wise(宏觀層面)【可以類比下 ConvNeXt】。
block-wise
首先,我們先從下列三個(gè)不同的視角進(jìn)行討論。
Patch-wise Features Projection
對比 ConvNets 和 ViTs 之間的相似性,兩個(gè)網(wǎng)絡(luò)都使用一個(gè)共同的基礎(chǔ)模塊將特征表示縮小到特定的尺寸。以往的方法大都沒有將圖像塊展平為具有線性層的順序輸入,因此作者采用具備大卷積核的投影層來提取 patch-wise 特征作為編碼器的輸入。
Volumetric Depth-wise Convolution with LKs
Swin transformer
的內(nèi)在特性之一是用于計(jì)算非局部 MSA 的滑動(dòng)窗口策略。總的來說,有兩種分層方法來計(jì)算 MSA:基于窗口的 MSA (W-MSA) 和 移動(dòng)窗口 MSA (SW-MSA)。這兩種方式都生成了跨層的全局感受野,并進(jìn)一步細(xì)化了非重疊窗口之間的特征對應(yīng)關(guān)系。
受深度卷積思想的啟發(fā),作者發(fā)現(xiàn)自注意力中的加權(quán)和方法與每通道卷積基礎(chǔ)之間的相似之處,其認(rèn)為使用 LK 大小的深度卷積可以在提取類似于 MSA 塊的特征時(shí)提供大的感受野。因此,本文建議采用 LK 大?。ɡ纾瑥?7 × 7 × 7 開始)通過深度卷積壓縮 Swin transformer 的窗口移動(dòng)特性。如此一來便可以保證每個(gè)卷積核通道與相應(yīng)的輸入通道進(jìn)行卷積運(yùn)算,使得輸出特征與輸入具備相同的通道維度。
Inverted Bottleneck with Depthwise Convolutional Scaling
Swin transformer
的另一個(gè)固有結(jié)構(gòu)是,它們被設(shè)計(jì)為 MLP 塊的隱藏層維度比輸入維度寬四倍,如下圖所示。有趣的是,這種設(shè)計(jì)與 ResNet 塊中的擴(kuò)張率相關(guān)。因此,我們利用 ResNet 塊中的類似設(shè)計(jì)并向上移動(dòng)深度卷積來計(jì)算特征。此外,通過引入了具有 1 × 1 × 1 卷積核大小的深度卷積縮放(DCS),以獨(dú)立地線性縮放每個(gè)通道特征。通過獨(dú)立擴(kuò)展和壓縮每個(gè)通道來豐富特征表示,可以最小化跨通道上下文產(chǎn)生的冗余信息,同時(shí)在每個(gè)階段增強(qiáng)了與下采樣塊的跨通道特征對應(yīng)。最后,通過使用 DCS,可以進(jìn)一步將模型復(fù)雜度降低 5%,并展示了與使用 MLP 模型的架構(gòu)相當(dāng)?shù)慕Y(jié)果。
layer-wise
介紹完微觀層面的設(shè)計(jì)思想,再讓我們從宏觀層面出發(fā),以另外三個(gè)嶄新的視角去理解作者的動(dòng)機(jī)。
Applying Residual Connections
從上圖左上角可以看到,標(biāo)準(zhǔn)的 3D U-Net 模型內(nèi)嵌的模塊為 2 個(gè) 3 x 3 x 3 卷積的堆疊,其展示了使用小卷積核提取具有增加通道的局部表示的樸素方法;而其右手邊的 SegResNet 則應(yīng)用了類似 3D 版本的瓶頸層,先降維再升維最后再接殘差表示;緊接著左下角的便是 Swin Transformer,其基于窗口注意力+MLP層的組合;最后右下角便是本文所提出的模塊,其在最后一個(gè)縮放層之后應(yīng)用輸入和提取特征之間的殘差連接。此外,在殘差求和前后并沒有應(yīng)用到任何的歸一化層和激活層。
Adapting Layer Normalization
我們知道,在卷積神經(jīng)網(wǎng)絡(luò)中,BN 是一種常用策略,它對卷積表示進(jìn)行歸一化以增強(qiáng)收斂性并減少過擬合。然而,之前的工作已經(jīng)證明 BN 會(huì)對模型的泛化能力產(chǎn)生不利影響。因此作者這里跟 ConvNeXt 一致,將 BN 替換為 LN。
Using GELU as the Activation Layer
ReLU 是個(gè)好東西,幾乎是現(xiàn)代 CNNs 模型的首選激活函數(shù)。作者在這里提倡使用 GELU,這是一種基于高斯誤差的線性變換單元,相對 ReLU 更加平滑,也是其中一種變體,解決 ReLU 因負(fù)梯度被硬截?cái)喽鴮?dǎo)致的神經(jīng)元失活問題。
方法
可以看出,整體的設(shè)計(jì)思路非常簡潔,框架也是一目了然,就一個(gè)標(biāo)準(zhǔn)的 3D U-Net 架構(gòu),由編碼器-解碼器組成,同時(shí)結(jié)合長跳躍連接操作幫助網(wǎng)絡(luò)更好的恢復(fù)空間細(xì)節(jié)的定位。也沒啥好分析的,下面就簡單拆開來看看,帶大家快速的過一遍就行。
DEPTH-WISE CONVOLUTION ENCODER
首先,輸入部分應(yīng)用一個(gè)大卷積核將原始圖像映射到一個(gè)低維空間分辨率的潛在空間特征表示,一來降低顯存參數(shù)量計(jì)算量等,二來可以增大網(wǎng)絡(luò)感受野,順便再處理成編碼器輸入所需的格式,總之好處多多,不過都是基操。
處理完之后就輸入到編碼器中進(jìn)行主要的特征提取,該編碼器共4個(gè)stage,也是標(biāo)準(zhǔn)的16倍下采樣。每個(gè)stage由多個(gè)不同的 3D UX-Net Block 構(gòu)成,具體長啥樣自己看下左圖一點(diǎn)就通。主要的特點(diǎn)就是主打輕量化和大感受野。
DECODER
編碼器中每個(gè)階段的多尺度輸出通過長跳躍連接鏈接到基于 ConvNet 的解碼器,并形成一個(gè)類似U形的網(wǎng)絡(luò)以用于下游的分割任務(wù)。這一塊就跟標(biāo)準(zhǔn)的 3D U-Net 幾乎沒啥兩樣了。
實(shí)驗(yàn)
消融實(shí)驗(yàn)
可視化效果
可以看到,與 GT 相比,3D UX-Net
器官和組織的形態(tài)得到了很好的保存。
與 SOTA 方法的對比
從上述表格可以看出,3D UX-Net 在所有分割任務(wù)中均展示出最佳性能,并且 Dice 分?jǐn)?shù)有了顯著提高(FeTA2021:0.870 到 0.874,F(xiàn)LARE2021:0.929 到 0.934)。
總結(jié)
本文為大家介紹了 3D UX-Net,這是第一個(gè)將分層 Transformer 的特征與用于醫(yī)學(xué)圖像分割的純 ConvNet 模塊相適應(yīng)的3D網(wǎng)絡(luò)架構(gòu)。具體地,本文重新設(shè)計(jì)了具有深度卷積的編碼器塊,以更低的代價(jià)實(shí)現(xiàn)與 Transformer 想媲美的能力。最后,通過在三個(gè)具有挑戰(zhàn)性的公共數(shù)據(jù)集上進(jìn)行的廣泛實(shí)驗(yàn)表明所提方法的高效性。
審核編輯 :李倩
-
神經(jīng)網(wǎng)絡(luò)
+關(guān)注
關(guān)注
42文章
4726瀏覽量
100339 -
圖像分割
+關(guān)注
關(guān)注
4文章
182瀏覽量
17953 -
數(shù)據(jù)集
+關(guān)注
關(guān)注
4文章
1199瀏覽量
24594
原文標(biāo)題:ICLR 2023 | 3D UX-Net:超強(qiáng)的醫(yī)學(xué)圖像分割新網(wǎng)絡(luò)
文章出處:【微信號(hào):CVer,微信公眾號(hào):CVer】歡迎添加關(guān)注!文章轉(zhuǎn)載請注明出處。
發(fā)布評論請先 登錄
相關(guān)推薦
評論