1. 摘要

UNet及其最新的擴(kuò)展如TransUNet是近年來(lái)領(lǐng)先的醫(yī)學(xué)圖像分割方法。然而，由于這些網(wǎng)絡(luò)參數(shù)多、計(jì)算復(fù)雜、使用速度慢，因此不能有效地用于即時(shí)應(yīng)用中的快速圖像分割。為此，我們提出了一種基于卷積多層感知器(MLP)的圖像分割網(wǎng)絡(luò)unext。我們?cè)O(shè)計(jì)了一種有效的UNeXt方法，即在前期采用卷積階段和在后期采用MLP階段。我們提出了一個(gè)標(biāo)記化的MLP塊，在該塊中，我們有效地標(biāo)記和投射卷積特征，并使用MLP來(lái)建模表示。

為了進(jìn)一步提高性能，我們建議在輸入mlp時(shí)shift輸入的channel，以便專注于學(xué)習(xí)局部依賴性。在潛在空間中使用標(biāo)記化的mlp減少了參數(shù)的數(shù)量和計(jì)算復(fù)雜度，同時(shí)能夠產(chǎn)生更好的表示，以幫助分割。該網(wǎng)絡(luò)還包括各級(jí)編碼器和解碼器之間的跳躍連接。測(cè)試結(jié)果表明，與目前最先進(jìn)的醫(yī)學(xué)圖像分割架構(gòu)相比，UNeXt的參數(shù)數(shù)量減少了72x，計(jì)算復(fù)雜度降低了68x，推理速度提高了10x，同時(shí)也獲得了更好的分割性能。

2. 網(wǎng)絡(luò)結(jié)構(gòu)

2.1 網(wǎng)絡(luò)設(shè)計(jì):

UNeXt是一個(gè)編碼器-解碼器體系結(jié)構(gòu)，有兩個(gè)階段:

1) 卷積階段

2) tokenized MLP階段。

輸入圖像通過(guò)編碼器，其中前3個(gè)塊是卷積，下2個(gè)是tokenized MLP塊。解碼器有2個(gè)tokenized MLP塊，后面跟著3個(gè)卷積塊。每個(gè)編碼器塊減少特征分辨率2倍，每個(gè)解碼器塊增加特征分辨率2。跳躍連接也被應(yīng)用在了編碼器和解碼器之間

作者減少了每個(gè)stage的通道數(shù)。

每個(gè)stage的通道數(shù)，對(duì)比標(biāo)準(zhǔn)的Unet：

UNeXt：32 64 128 160 256

UNet：64 128 256 512 1024

在這里面就減少了很多的參數(shù)量

2.2 卷積階段

有三個(gè)conv block，每個(gè)block都有一個(gè)卷積層（傳統(tǒng)Unet是兩個(gè)）、批量歸一化層和ReLU激活。我們使用的內(nèi)核大小為3×3, stride為1,padding為1。編碼器的conv塊使用帶有池窗口2×2的max-pooling層，而解碼器的conv塊使用雙線性插值層對(duì)特征圖進(jìn)行上采樣。我們使用雙線性插值而不是轉(zhuǎn)置卷積，因?yàn)檗D(zhuǎn)置卷積基本上是可學(xué)習(xí)的上采樣，會(huì)導(dǎo)致產(chǎn)生更多可學(xué)習(xí)的參數(shù)

2.3 Shifted MLP

在shifted MLP中，在tokenize之前，我們首先移動(dòng)conv features通道的軸線。這有助于MLP只關(guān)注conv特征的某些位置，從而誘導(dǎo)塊的位置。這里的直覺(jué)與Swin transformer類似，在swin中引入基于窗口的注意，以向完全全局的模型添加更多的局域性。由于Tokenized MLP塊有2個(gè)mlp，我們?cè)谝粋€(gè)塊中跨越寬度移動(dòng)特征，在另一個(gè)塊中跨越高度移動(dòng)特征，就像軸向注意力中一樣。我們對(duì)這些特征做了h個(gè)劃分，并根據(jù)指定的軸通過(guò)j個(gè)位置移動(dòng)它們。這有助于我們創(chuàng)建隨機(jī)窗口，引入沿軸線的局部性。

Shift操作

圖中灰色是特征塊的位置，白色是移動(dòng)之后的padding。

2.4 Tokenized MLP階段

image-20220402001733482

在Tokenized MLP塊中，我們首先shift features并將它們投射到token中。為了進(jìn)行token化，我們首先使用3x3conv把特征投射到E維，其中E是embadding維度(token的數(shù)量)，它是一個(gè)超參數(shù)。然后我們將這些token傳遞給一個(gè)shifted MLP(跨越width)。接下來(lái)，特征通過(guò) DW-Conv傳遞。然后我們使用GELU激活層。然后，我們通過(guò)另一個(gè)shifted MLP(跨越height)傳遞特征，該mlp把特征的尺寸從H轉(zhuǎn)換為了O。我們?cè)谶@里使用一個(gè)殘差連接，并將原始標(biāo)記添加為殘差。然后我們利用layer norm（LN），并將輸出特征傳遞到下一個(gè)塊。LN比BN更可取，因?yàn)樗幸饬x的是沿著token進(jìn)行規(guī)范化，而不是在Tokenized MLP塊的整個(gè)批處理中進(jìn)行規(guī)范化。

我們?cè)谶@個(gè)塊中使用DWConv有兩個(gè)原因:

1)它有助于編碼MLP特征的位置信息。從中可以看出，在一個(gè)MLP塊中Conv層已經(jīng)足夠?qū)ξ恢眯畔⑦M(jìn)行編碼，并且實(shí)際性能優(yōu)于標(biāo)準(zhǔn)的位置編碼技術(shù)。當(dāng)測(cè)試或者訓(xùn)練分辨率不相同時(shí)，像ViT中的位置編碼技術(shù)需要插值，這通常會(huì)導(dǎo)致性能下降。

2)DWConv使用更少的參數(shù)，因此提高了效率。

Tokenized block的計(jì)算流程

所有這些計(jì)算都是在嵌入維數(shù)h上執(zhí)行的，這個(gè)維數(shù)明顯小于特征的維數(shù) (H/N)×(H/N) ,N是關(guān)于降維的2的因子。在我們的實(shí)驗(yàn)中，除非另有說(shuō)明，否則我們使用768。這種設(shè)計(jì)tokenized MLP block的方法有助于編碼有意義的特征信息，而不會(huì)對(duì)計(jì)算或參數(shù)貢獻(xiàn)太多。

3.實(shí)驗(yàn)結(jié)果

在ISIC和BUSI數(shù)據(jù)集進(jìn)行了實(shí)驗(yàn)

在ISIC數(shù)據(jù)集的對(duì)比

4. 個(gè)人感悟

首先每個(gè)convolutional階段只有一個(gè)卷積層，極大的減少了運(yùn)算量，是答主第一次見(jiàn)了。

其次是把MLP的模塊引入了Unet，算是很新穎了。

在Tokenized MLP block中使用DW- CONV，讓人眼前一亮。