來源丨 GiantPandaCV

• 前言
• 方法概述
• UNeXt 架構(gòu)
• TokMLP 設計思路
• 實驗部分
• 一些理解和總結(jié)
• 參考鏈接

最近 MICCAI 2022 的論文集開放下載了，地址：https://link.springer.com/book/10.1007/978-3-031-16443-9 ，每個部分的內(nèi)容如下所示：

Part I: Brain development and atlases; DWI and tractography; functional brain networks; neuroimaging; heart and lung imaging; dermatology;

Part II: Computational (integrative) pathology; computational anatomy and physiology; ophthalmology; fetal imaging;

Part III: Breast imaging; colonoscopy; computer aided diagnosis;

Part IV: Microscopic image analysis; positron emission tomography; ultrasound imaging; video data analysis; image segmentation I;

Part V: Image segmentation II; integration of imaging with non-imaging biomarkers;

Part VI: Image registration; image reconstruction;

Part VII: Image-Guided interventions and surgery; outcome and disease prediction; surgical data science; surgical planning and simulation; machine learning – domain adaptation and generalization;

Part VIII: Machine learning – weakly-supervised learning; machine learning – model interpretation; machine learning – uncertainty; machine learning theory and methodologies.

其中關(guān)于分割有兩個部分，Image segmentation I 在 Part IV, 而 Image segmentation II 在 Part V。計劃對其中開放源代碼和典型的方法注意解讀，這次要分享的論文是其中的 UNeXt: MLP-based Rapid Medical Image Segmentation Network，arXiv 鏈接：https://arxiv.org/abs/2203.04967 。

隨著醫(yī)學圖像的解決方案變得越來越適用，我們更需要關(guān)注使深度網(wǎng)絡輕量級、快速且高效的方法。具有高推理速度的輕量級網(wǎng)絡可以被部署在手機等設備上，例如 POCUS（point-of-care ultrasound）被用于檢測和診斷皮膚狀況。這就是 UNeXt 的動機。

之前我們解讀過基于 Transformer 的 U-Net 變體，近年來一直是領(lǐng)先的醫(yī)學圖像分割方法，但是參數(shù)量往往不樂觀，計算復雜，推理緩慢。這篇文章提出了基于卷積多層感知器（MLP）改進 U 型架構(gòu)的方法，可以用于圖像分割。設計了一個 tokenized MLP 塊有效地標記和投影卷積特征，使用 MLPs 來建模表示。這個結(jié)構(gòu)被應用到 U 型架構(gòu)的下兩層中（這里我們假設縱向一共五層）。文章中提到，為了進一步提高性能，建議在輸入到 MLP 的過程中改變輸入的通道，以便專注于學習局部依賴關(guān)系特征。還有額外的設計就是跳躍連接了，并不是我們主要關(guān)注的地方。最終，UNeXt 將參數(shù)數(shù)量減少了 72 倍，計算復雜度降低了 68 倍，推理速度提高了 10 倍，同時還獲得了更好的分割性能，如下圖所示。

UNeXt 的設計如下圖所示?？v向來看，一共有兩個階段，普通的卷積和 Tokenized MLP 階段。其中，編碼器和****分別設計兩個 Tokenized MLP 塊。每個編碼器將分辨率降低兩倍，****工作相反，還有跳躍連接結(jié)構(gòu)。每個塊的通道數(shù)（C1-C5）被設計成超參數(shù)為了找到不掉點情況下最小參數(shù)量的網(wǎng)絡，對于使用 UNeXt 架構(gòu)的實驗，遵循 C1 = 32、C2 = 64、C3 = 128、C4 = 160 和 C5 = 256。

關(guān)于 Convolutional Stage 我們不做過多介紹了，在這一部分重點專注 Tokenized MLP Stage。從上一部分的圖中，可以看到 Shifted MLP 這一操作，其實思路類似于 Swin transformer，引入基于窗口的注意力機制，向全局模型中添加更多的局域性。下圖的意思是，Tokenized MLP 塊有 2 個 MLP，在一個 MLP 中跨越寬度移動特征，在另一個 MLP 中跨越高度移動特征，也就是說，特征在高度和寬度上依次移位。論文中是這么說的：“我們將特征分成 h 個不同的分區(qū)，并根據(jù)指定的軸線將它們移到 j=5 的位置”。其實就是創(chuàng)建了隨機窗口，這個圖可以理解為灰色是特征塊的位置，白色是移動之后的 padding。

解釋過 Shifted MLP 后，我們再看另一部分：tokenized MLP block。首先，需要把特征轉(zhuǎn)換為 tokens（可以理解為 Patch Embedding 的過程）。為了實現(xiàn) tokenized 化，使用 kernel size 為 3 的卷積，并將通道的數(shù)量改為 E，E 是 embadding 嵌入維度（ token 的數(shù)量），也是一個超參數(shù)。然后把這些 token 送到上面提到的第一個跨越寬度的 MLP 中。

這里會產(chǎn)生了一個疑問，關(guān)于 kernel size 為 3 的卷積，使用的是什么樣的卷積層？答：這里還是普通的卷積，文章中提到了 DWConv（DepthWise Conv），是后面的特征通過 DW-Conv 傳遞。使用 DWConv 有兩個原因：（1）它有助于對 MLP 特征的位置信息進行編碼。MLP 塊中的卷積層足以編碼位置信息，它實際上比標準的位置編碼表現(xiàn)得更好。像 ViT 中的位置編碼技術(shù)，當測試和訓練的分辨率不一樣時，需要進行插值，往往會導致性能下降。（2）DWConv 使用的參數(shù)數(shù)量較少。

這時我們得到了 DW-Conv 傳遞過來的特征，然后使用 GELU 完成激活。接下來，通過另一個 MLP（跨越height）傳遞特征，該 MLP 把進一步改變了特征尺寸。在這里還使用一個殘差連接，將原始 token 添加為殘差。然后我們利用 Layer Norm（LN），將輸出特征傳遞到下一個塊。LN 比 BN 更可取，因為它是沿著 token 進行規(guī)范化，而不是在 Tokenized MLP 塊的整個批處理中進行規(guī)范化。上面這些就是一個 tokenized MLP block 的設計思路。

此外，文章中給出了 tokenized MLP block 涉及的計算公式：

其中 T 表示 tokens，H 表示高度，W 表示寬度。值得注意的是，所有這些計算都是在 embedding 維度 H 上進行的，它明顯小于特征圖的維度 HN×HN，其中 N 取決于 block 大小。在下面的實驗部分，文章將 H 設置為 768。

實驗在 ISIC 和 BUSI 數(shù)據(jù)集上進行，可以看到，在 GLOPs、性能和推理時間都上表現(xiàn)不錯。

下面是可視化和消融實驗的部分?？梢暬瘓D可以發(fā)現(xiàn)，UNeXt 處理的更加圓滑和接近真實標簽。

消融實驗可以發(fā)現(xiàn)，從原始的 UNet 開始，然后只是減少過濾器的數(shù)量，發(fā)現(xiàn)性能下降，但參數(shù)并沒有減少太多。接下來，僅使用 3 層深度架構(gòu)，既 UNeXt 的 Conv 階段。顯著減少了參數(shù)的數(shù)量和復雜性，但性能降低了 4%。加入 tokenized MLP block 后，它顯著提高了性能，同時將復雜度和參數(shù)量是一個最小值。接下來，我們將 DWConv 添加到 positional embedding，性能又提高了。接下來，在 MLP 中添加  Shifted 操作，表明在標記化之前移位特征可以提高性能，但是不會增加任何參數(shù)或復雜性。注意：Shifted MLP 不會增加 GLOPs。

在這項工作中，提出了一種新的深度網(wǎng)絡架構(gòu) UNeXt，用于醫(yī)療圖像分割，專注于參數(shù)量的減小。UNeXt 是一種基于卷積和 MLP 的架構(gòu)，其中有一個初始的 Conv 階段，然后是深層空間中的 MLP。具體來說，提出了一個帶有移位 MLP 的標記化 MLP 塊。在多個數(shù)據(jù)集上驗證了 UNeXt，實現(xiàn)了更快的推理、更低的復雜性和更少的參數(shù)數(shù)量，同時還實現(xiàn)了最先進的性能。

我在讀這篇論文的時候，直接注意到了它用的數(shù)據(jù)集。我認為 UNeXt 可能只適用于這種簡單的醫(yī)學圖像分割任務，類似的有 Optic Disc and Cup Seg，對于更復雜的，比如血管，軟骨，Liver Tumor，kidney Seg 這些，可能效果達不到這么好，因為運算量被極大的減少了，每個 convolutional 階段只有一個卷積層。MLP 魔改 U-Net 也算是一個嘗試，在 Tokenized MLP block 中加入 DWConv 也是很合理的設計。

1、https://jeya-maria-jose.github.io/UNext-web/2、https://arxiv.org/abs/2203.049673、https://github.com/jeya-maria-jose/UNeXt-pytorch

本文僅做學術(shù)分享，如有侵權(quán)，請聯(lián)系刪文。