UNeXt: MLP-based Rapid Medical ImageSegmentation Network

最新推荐文章于 2024-08-14 20:41:00 发布

shchojj

最新推荐文章于 2024-08-14 20:41:00 发布

阅读量552

点赞数 1

分类专栏： segmentation 文章标签：深度学习

原文链接：https://arxiv.org/abs/2203.04967

版权

segmentation 专栏收录该内容

32 篇文章 16 订阅

订阅专栏

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编写的损失函数是