MICCAI 2022：基于 MLP 的快速醫(yī)學(xué)圖像分割網(wǎng)絡(luò) UNeXt

來源丨 GiantPandaCV

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

最近 MICCAI 2022 的論文集開放下載了，地址：https://link.springer.com/book/10.1007/978-3-031-16443-9 ，每個(gè)部分的內(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)于分割有兩個(gè)部分，Image segmentation I 在 Part IV, 而 Image segmentation II 在 Part V。計(jì)劃對(duì)其中開放源代碼和典型的方法注意解讀，這次要分享的論文是其中的 UNeXt: MLP-based Rapid Medical Image Segmentation Network，arXiv 鏈接：https://arxiv.org/abs/2203.04967 。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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