UNeXt：第一个基于卷积和MLP的快速医学图像分割网络

会员服务 ·

UNeXt：第一个基于卷积和MLP的快速医学图像分割网络

2022 年 4 月 4 日 CVer

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

作者：Wildeur | 已授权转载（源：知乎）编辑：CVer

https://zhuanlan.zhihu.com/p/491798644

UNeXt: MLP-based Rapid Medical Image Segmentation Network

论文：https://arxiv.org/abs/2203.04967

代码（基于PyTorch，已开源）：

https://github.com/jeya-maria-jose/UNeXt-pytorch

1. 摘要

UNet及其最新的扩展如TransUNet是近年来领先的医学图像分割方法。然而，由于这些网络参数多、计算复杂、使用速度慢，因此不能有效地用于即时应用中的快速图像分割。为此，我们提出了一种基于卷积多层感知器(MLP)的图像分割网络unext。我们设计了一种有效的UNeXt方法，即在前期采用卷积阶段和在后期采用MLP阶段。我们提出了一个标记化的MLP块，在该块中，我们有效地标记和投射卷积特征，并使用MLP来建模表示。

为了进一步提高性能，我们建议在输入mlp时shift输入的channel，以便专注于学习局部依赖性。在潜在空间中使用标记化的mlp减少了参数的数量和计算复杂度，同时能够产生更好的表示，以帮助分割。该网络还包括各级编码器和解码器之间的跳跃连接。测试结果表明，与目前最先进的医学图像分割架构相比，UNeXt的参数数量减少了72x，计算复杂度降低了68x，推理速度提高了10x，同时也获得了更好的分割性能。

2. 网络结构

2.1 网络设计:

UNeXt是一个编码器-解码器体系结构，有两个阶段:

1) 卷积阶段

2) tokenized MLP阶段。

输入图像通过编码器，其中前3个块是卷积，下2个是tokenized MLP块。解码器有2个tokenized MLP块，后面跟着3个卷积块。每个编码器块减少特征分辨率2倍，每个解码器块增加特征分辨率2。跳跃连接也被应用在了编码器和解码器之间

作者减少了每个stage的通道数。

每个stage的通道数，对比标准的Unet：

UNeXt：32 64 128 160 256
UNet：64 128 256 512 1024

在这里面就减少了很多的参数量

2.2 卷积阶段

有三个conv block，每个block都有一个卷积层（传统Unet是两个）、批量归一化层和ReLU激活。我们使用的内核大小为3×3, stride为1,padding为1。编码器的conv块使用带有池窗口2×2的max-pooling层，而解码器的conv块使用双线性插值层对特征图进行上采样。我们使用双线性插值而不是转置卷积，因为转置卷积基本上是可学习的上采样，会导致产生更多可学习的参数

2.3 Shifted MLP

在shifted MLP中，在tokenize之前，我们首先移动conv features通道的轴线。这有助于MLP只关注conv特征的某些位置，从而诱导块的位置。这里的直觉与Swin transformer类似，在swin中引入基于窗口的注意，以向完全全局的模型添加更多的局域性。由于Tokenized MLP块有2个mlp，我们在一个块中跨越宽度移动特征，在另一个块中跨越高度移动特征，就像轴向注意力中一样。我们对这些特征做了h个划分，并根据指定的轴通过j个位置移动它们。这有助于我们创建随机窗口，引入沿轴线的局部性。

图中灰色是特征块的位置，白色是移动之后的padding。

2.4 Tokenized MLP阶段

在Tokenized MLP块中，我们首先shift features并将它们投射到token中。为了进行token化，我们首先使用3x3conv把特征投射到E维，其中E是embadding维度(token的数量)，它是一个超参数。然后我们将这些token传递给一个shifted MLP(跨越width)。接下来，特征通过 DW-Conv传递。然后我们使用GELU激活层。然后，我们通过另一个shifted MLP(跨越height)传递特征，该mlp把特征的尺寸从H转换为了O。我们在这里使用一个残差连接，并将原始标记添加为残差。然后我们利用layer norm（LN），并将输出特征传递到下一个块。LN比BN更可取，因为它更有意义的是沿着token进行规范化，而不是在Tokenized MLP块的整个批处理中进行规范化。

我们在这个块中使用DWConv有两个原因:

1)它有助于编码MLP特征的位置信息。从中可以看出，在一个MLP块中Conv层已经足够对位置信息进行编码，并且实际性能优于标准的位置编码技术。当测试或者训练分辨率不相同时，像ViT中的位置编码技术需要插值，这通常会导致性能下降。

2)DWConv使用更少的参数，因此提高了效率。

Tokenized block的计算流程

所有这些计算都是在嵌入维数h上执行的，这个维数明显小于特征的维数 (H/N)×(H/N) ,N是关于降维的2的因子。在我们的实验中，除非另有说明，否则我们使用768。这种设计tokenized MLP block的方法有助于编码有意义的特征信息，而不会对计算或参数贡献太多。

3.实验结果

在ISIC和BUSI数据集进行了实验

在ISIC数据集的对比

4. 个人感悟

首先每个convolutional阶段只有一个卷积层，极大的减少了运算量，是答主第一次见了。

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

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

UNeXT 论文和代码下载

后台回复：UNeXT，即可下载上述论文和代码

后台回复：ICCV2021，即可下载ICCV 2021论文和代码开源的论文合集

后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集

重磅！医学图像 交流群成立

扫描下方二维码，或者添加微信：CVer6666，即可添加CVer小助手微信，便可申请加入CVer-医疗影像微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer等。

一定要备注：研究方向+地点+学校/公司+昵称（如医疗影像+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群