天天看點

UNeXt:基于 MLP 的快速醫學圖像分割網絡

作者:deephub

UNeXt是約翰霍普金斯大學在2022年釋出的論文。它在早期階段使用卷積,在潛在空間階段使用 MLP。通過一個标記化的 MLP 塊來标記和投影卷積特征,并使用 MLP 對表示進行模組化。對輸入通道進行移位,可以專注于學習局部依賴性。

UNeXt

UNeXt:基于 MLP 的快速醫學圖像分割網絡

基本架構

UNeXt 是一種編碼器-解碼器架構,具有兩個階段:

1、卷積階段,2、标記化 MLP 階段。

輸入圖像通過編碼器,前 3 個塊是卷積塊,接下來的 2 個是标記化 MLP 塊。

解碼器有 2 個标記化 MLP 塊,後跟 3 個卷積塊。

每個編碼器塊使用具有視窗 2×2 的最大池化層将特征分辨率降低 2,每個解碼器塊使用雙線性插值将特征分辨率增加 2。編碼器和解碼器之間也包含跳過連接配接。 每個塊的通道數是一個超參數,表示為 C1 到 C5。 在實驗中,除非另有說明,否則 C1=32、C2=64、C3=128、C4=160 和 C5=256。

每個卷積塊時标準的一個卷積層、一個批量歸一化層和 ReLU 激活層。 核心大小為 3×3,步長為 1。

帶位移的 MLP

UNeXt:基于 MLP 的快速醫學圖像分割網絡

conv特征的通道軸線在标記(Tokenized)之前首先移位。這有助于MLP隻關注conv特征的某些位置,進而誘導塊的局部性。論文作者說,這裡與Swin Transformer類似。由于Tokenized MLP塊有2個MLP,是以特征在一個塊中跨寬度移動,在另一個塊中跨高度移動,就像Axial-DeepLab中的軸向注意力一樣。這樣特征被分割到h個不同的分區,并根據指定的軸移動j=5個位置。

标記化(Tokenized) MLP階段

UNeXt:基于 MLP 的快速醫學圖像分割網絡

首先使用大小為 3 的核心将通道數更改為嵌入次元 E(标記數)。然後将這些标記令牌傳遞給一個帶移位的MLP(跨寬度),其中包含MLP的隐藏次元,預設H=768。

接下來,使用深度卷積層(DWConv)。它有助于對位置資訊進行編碼,像SegFormer中所建議的,當訓練/測試分辨率不同時,它比ViT具有更好的性能。并且它使用更少的參數,可以提高了效率。

激活函數使用GELU,因為在ViT和BERT在使用GELU的情況下表現更好。

特征通過另一個移位的MLP(跨高度)傳遞,該MLP将次元從H轉換為O。

最後還是用了殘差連接配接将原始标記令牌添加到殘差。然後使用層歸一化(LN),将輸出特征傳遞給下一個塊。

損失函數

使用二進制交叉熵(BCE)和dice 損失的組合:

UNeXt:基于 MLP 的快速醫學圖像分割網絡

結果展示

SOTA對比

UNeXt:基于 MLP 的快速醫學圖像分割網絡

UNeXt獲得了比所有基線更好的分割性能,計算量比第二的TransUNet少得多。UNeXt在計算複雜度方面明顯優于所有其他網絡。

UNeXt:基于 MLP 的快速醫學圖像分割網絡

swing - unet(圖中未顯示)有41.35 M個參數,計算也很複雜有11.46 GFLOPs。

作者還實驗了MLP-Mixer作為編碼器和普通卷積解碼器,它隻有大約11M個參數,但是分割的性能不是最優的。

定性結果

UNeXt:基于 MLP 的快速醫學圖像分割網絡

與其他方法相比,UNeXt産生了具有競争力的分割預測。

消融實驗

UNeXt:基于 MLP 的快速醫學圖像分割網絡

當深度減小,僅使用3級架構,也就是說隻使用Conv階段時,參數數量和複雜度顯著減少,但性能下降4%。當使用标記化的MLP塊時,它可以顯着提高性能。

UNeXt:基于 MLP 的快速醫學圖像分割網絡

增加通道(UNeXt-L)進一步提高了性能,同時增加了計算開銷。減少通道(UNeXt-S)會降低性能(降低幅度并不大),但我們得到了一個非常輕量級的模型。

論文:https://arxiv.org/abs/2203.04967

源代碼:https://github.com/jeya-maria-jose/UNeXt-pytorch

本文作者:Sik-Ho Tsang

繼續閱讀