UNeXt是約翰霍普金斯大學在2022年釋出的論文。它在早期階段使用卷積，在潛在空間階段使用 MLP。通過一個标記化的 MLP 塊來标記和投影卷積特征，并使用 MLP 對表示進行模組化。對輸入通道進行移位，可以專注于學習局部依賴性。
UNeXt
基本架構
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
conv特征的通道軸線在标記（Tokenized)之前首先移位。這有助于MLP隻關注conv特征的某些位置，進而誘導塊的局部性。論文作者說，這裡與Swin Transformer類似。由于Tokenized MLP塊有2個MLP，是以特征在一個塊中跨寬度移動，在另一個塊中跨高度移動，就像Axial-DeepLab中的軸向注意力一樣。這樣特征被分割到h個不同的分區，并根據指定的軸移動j=5個位置。
标記化（Tokenized） MLP階段
首先使用大小為 3 的核心将通道數更改為嵌入次元 E（标記數）。然後将這些标記令牌傳遞給一個帶移位的MLP(跨寬度)，其中包含MLP的隐藏次元，預設H=768。
接下來，使用深度卷積層(DWConv)。它有助于對位置資訊進行編碼，像SegFormer中所建議的，當訓練/測試分辨率不同時，它比ViT具有更好的性能。并且它使用更少的參數，可以提高了效率。
激活函數使用GELU，因為在ViT和BERT在使用GELU的情況下表現更好。
特征通過另一個移位的MLP(跨高度)傳遞，該MLP将次元從H轉換為O。
最後還是用了殘差連接配接将原始标記令牌添加到殘差。然後使用層歸一化(LN)，将輸出特征傳遞給下一個塊。
損失函數
使用二進制交叉熵(BCE)和dice 損失的組合:
結果展示
SOTA對比
UNeXt獲得了比所有基線更好的分割性能，計算量比第二的TransUNet少得多。UNeXt在計算複雜度方面明顯優于所有其他網絡。
swing - unet(圖中未顯示)有41.35 M個參數，計算也很複雜有11.46 GFLOPs。
作者還實驗了MLP-Mixer作為編碼器和普通卷積解碼器，它隻有大約11M個參數，但是分割的性能不是最優的。
定性結果
與其他方法相比，UNeXt産生了具有競争力的分割預測。
消融實驗
當深度減小，僅使用3級架構，也就是說隻使用Conv階段時，參數數量和複雜度顯著減少，但性能下降4%。當使用标記化的MLP塊時，它可以顯着提高性能。
增加通道(UNeXt-L)進一步提高了性能，同時增加了計算開銷。減少通道(UNeXt-S)會降低性能(降低幅度并不大)，但我們得到了一個非常輕量級的模型。
論文：https://arxiv.org/abs/2203.04967
源代碼：https://github.com/jeya-maria-jose/UNeXt-pytorch
本文作者：Sik-Ho Tsang