UNeXt: MLP-based Rapid Medical ImageSegmentation Network

UneXt关注的是移动端的设备，比如point-of-care ultrasound (POCUS)、

 Phone-based skin lesion detection and identification application
、Magnetic resonance imaging (MRI)

 

 

 因为专注的是移动端，因此效率和准确度是UneXt关注的重点。效率的设计主要是减少网络参数，具体的设计模式在后面再介绍。作者描述UNeXt参数减少72倍，计算复杂度减少68倍，具体的参数是1.47m， GFLOPs是0.57。

论文的贡献点： 

1、第一个convolutional MLP-base 图像分割网络（convolutional multilayer perception）

2、用于在潜在空间学习较好表示的tokenized MLP block with axial shifts

3、在不降低分割性能的前提下，提升了分割效率。

 网络的encoder-decoder含两个stage，在decoder和encoder上有3个convolutional block和2个 tokenized MLP blocks，分辨率都是以2倍的方式进行改变，skip connections 不用说肯定也有的。通道数的话，论文中使用的是C1=32, C2=64, C3=128, C4=160，C5=256，前面三个都是成倍变化的，因为是卷积么，论文中提filters数减少了，从而参数和计算都减少了。

一、convolutional stage：

一个convolutional block就是，也就是convolution layer,  batch normalization layer and ReLU。kernel size = 3x3, stride =1, padding =1。encoder中上采样环节不在使用transpose convolutional ，而使用bilinear interpolation layer，这样的话反卷积中的参数也省掉了，max-pooling layer的pool window = 2x2

二、Tokenized MLP stage

  

首先需要介绍shifted MLP, 在tokenizing 之前，先沿着channel方向shift axis。这样可以帮助MLP值关注某些location 从而在block中引入locality信息。这里类似于swin transformer 中的window-based attention，即林如locality到global模型。

 tokenized MLP中包含两个 shift  MLP，即width方向和height方向，这个和axial-attention原理一样，也是就是将channel沿着width、height两个方向shift。将特征h个不同的部分，并沿着指定的axis移动到对应位置j。这有助于我们创建随机窗口，沿着axis方向引入locality信息。

shifted mlp中描述的就是将feature进行shift，并投影到tokens，投影使用kernel size =3并将channel改变为E，E是一个embedding dimension，即tokens个数，这是一个超参。

然后将tokens 通过shifted MPL(across width)，这里面的MLP隐藏参数H也是一个超参。

 

然后通过一个depth wise convolution layer(DWConv)。

DWConv有两个好处：

1、帮助MLP features编码positional 信息。实际上是一个MLP block中的卷积层，对于位置编码的能力要比标准的位置编码技术要强，比如ViT中的ones位置编码技术，如果测试和训练的分辨率不一致时，会导致性能下降。

2、另外DWConv使用更少的参数，因而效率得到提升。

DWConv之后使用了GELU激活层，之所以使用GELU而不是RELU是因为它更平滑且性能更好。而且在最近的ViT和BERT结果也是证明GELU提升了性能。

随后通过一个shifted-MLP(across height)将维度从H转换到O

 

layer normalization(LN) 作为归一化层，之所以使用LN而不是BN，是因为沿着tokens方向归一化比验证batch方向归一化更有意义

 residual connection追加的是原始的tokens

  shifted MLP Width

 MLP+ GELU+Encoding

 shifted MLP Height

 layer normalization(MLP(GELU)+tokenize)

 所有这些计算都是在嵌入维度H上执行的，它明显小于特征映射H/N × H/N的维度，其中N是与block相关的2的整数倍。论文中H=768，这里设计的Tokenized MLP block可以编码更多有意义的特征信息而不增加额外的计算量和参数。

论文中使用的数据集是International Skin Imaging Collaboration (ISIC 2018) 皮肤病数据集 and Breast UltraSound Images (BUSI)乳腺癌超声构建benchmark。ISIC 2018有2594张图像，BUSI有647张图像，并重采样到256x256.

代码是pytorch编写的损失函数是

 binary cross entropy(BEC)和Dice loss

优化器时Adam，learning rate = 0.0001，momentum =0.9，并且使用minimum learning rate>0.00001的余弦退火学习率策略。

batch size = 8

epochs = 400

使用80:20的随机策略划分数据集，执行三次训练，并输出均值和方差。

性能比较

  原本是想用Swin-UNet做对比实验，但是在数据集比较少的时候，Swin-UNet性能就很差。

 消融实验

Reduced UNet就是减少了filters数，来减少参数和复杂度，但是分割性能没有下降。

Conv Stage减少深度，并且仅使用3层卷积(UNeXt中的)，虽然参数和复杂度减少了很多，但性能下降了4%

Conv Stage+ Tok-MLP w/o PE就是在Conv Stage中加入tokenized MLP block，可以发现性能提升了，但是参数和复杂度只是少量增加。

Conv Stage+ Tok-MLP + PE就是在上一步增加了DWConv 的 positional embedding method

Conv Stage+ Shifted Tok-MLP(W) + PE就是将MLPs的shifting操作方向改为W

Conv Stage+ Shifted Tok-MLP(H) + PE就是将MLPs的shifting操作方向改为W

Conv Stage+ Shifted Tok-MLP(H+W) + PE就是将MLPs的shifting操作方向改为H+W

sihft操作不会增加计算复杂度。所有的实验都是在ISIC数据集上进行单折实验。

通道数分析

channels数是影响网络参数、复杂度、性能主要超参。实验都是在ISIC一折上进行测试，通道数增加，性能肯定提升了，但是网络参数和性能也下降了。

 MLP-Mixer差异

网络使用MLP-Mixer关注于channel 和token的mixing来学习更好的表达，使用其他的MLP策略不但性能下降而且参数也增多了。