UNeXT

论文链接：https://link.springer.com/chapter/10.1007/978-3-031-16443-9_3

源码链接：https://github.com/jeya-maria-jose/UNeXt-pytorch

论文摘要

摘要近年来，UNet及其最新扩展(如TransUNet)已成为医学图像分割的主要方法。然而，这些网络不能有效地用于即时护理应用中的快速图像分割，因为它们参数量大，计算复杂，使用速度慢。为此，我们提出了一种基于卷积多层感知器(MLP)的图像分割网络UNeXt。我们有效地设计了UNeXt，早期是卷积阶段，潜伏阶段（latent stage）是MLP。我们提出了一个标记化的MLP块，其中我们有效地标记和投影卷积特征，并使用MLP来建模表示。为了进一步提高性能，我们建议在输入到mlp时改变输入的通道，以便专注于学习局部依赖。在潜在空间中使用标记化的mlp减少了参数的数量和计算复杂性，同时能够产生更好的表示来帮助分割。该网络还包括各级编码器和解码器之间的跳越连接（skip connection）。我们在多个医学图像分割数据集上测试了UNeXt，结果表明，我们将参数数量减少了72倍，计算复杂度降低了68倍，推理速度提高了10倍，同时还获得了比最先进的医学图像分割架构更好的分割性能。

结构详解

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

1)卷积阶段（convolutional stage)

2)标记化MLP阶段( Tokenized MLP stage)

-输入图像通过编码器传递，其中前3个块是卷积块，接下来的2个是标记化MLP块。解码器有2个标记化MLP块和3个卷积块。每个编码器块将特征分辨率降低2，每个解码器块将特征分辨率提高2。

-在编码器和解码器之间存在残差链接（Skip connections）

-每个块上的通道数量是一个超参数，表示为C1到C5

-对于使用UNeXt架构，我们遵循C1 = 32, C2 = 64, C3 = 128, C4 = 160，和除非另有说明，C5 = 256。请注意，这些数字实际上小于UNet的过滤器数量及其变体的数量，以减少参数和计算量。

卷积阶段Convolutional Stage:

每个卷积块都配有卷积层、批规范化层（batch normalization）和ReLU激活。我们使用的内核大小为3 × 3，步幅为1，填充为1。编码器中的conv块使用池窗口为2 × 2的最大池化层，而解码器中的conv块由双线性插值层组成，对特征映射进行上采样。我们使用双线性插值（bilinear interpolation）代替转置卷积，因为转置卷积基本上是可学习的上采样，可学习的参数更多。

*Shifted MLP:

在shifted MLP中，我们首先移位conv的通道轴标记之前的特性。这有助于MLP只关注conv特征的特定位置，从而诱导块的局部性。这里的直觉（intuition）是类似于Swin transformer，其中引入了基于窗口的注意力（ window-based attention），为一个完全全局的模型添加更多的局部性Tokenized MLP块有2个MLP，我们在其中一个上面通过宽度移动特征，另一个通过高度移动特征，就如 axial-attention一样。我们把特征分成h个不同的分区，并根据指定的轴按j个位置移动它们。这有助于我们创建随机窗口，并沿轴引入局部性。

Tokenized MLP Stage:

在标记化的MLP块中，我们首先转移特性并将它们投射到标记（token）中。为了标记化，我们首先使用内核大小为3，并将通道数量更改为E，其中E是嵌入维度（ embedding dimension）(token数量)，这是一个超参数。然后，我们将这些标记传递给一个shifted MLP(跨宽度)，其中MLP的隐藏维度是一个超参数h。接下来，特征通过深度卷积层(DWConv)传递。

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

1)它有助于编码MLP特征的位置信息。MLP块中的Conv层足以对位置信息进行编码，并且它实际上比标准的位置编码技术性能更好。像ViT中的位置编码技术需要在测试和训练分辨率不相同时进行插值，这通常会导致性能降低。

2) DWConv使用较少的参数，从而提高了效率。然后我们使用GELU激活层。GELU是RELU更流畅的替代品，并且被发现性能更好。此外，最近的架构如ViT[10]和BERT[9]已经成功地使用GELU获得了改进的结果

然后，我们将特征传递给另一个shifted MLP(跨高度)，它将维度从H转换为o。我们在这里使用残差连接（residual connection），并将原始标记作为残差添加。然后我们应用 层归一化（layer normalization）(LN)并将输出特征传递给下一个块。LN优于BN，因为在标记化MLP块中沿着标记进行归一化比在批处理中进行归一化更有意义。

在Tokenized MLP的计算可以总结为

其中T表示token，H表示高度，W表示宽度，DW Conv表示深度卷积，LN表示层归一化。请注意，所有这些计算都是在嵌入维度（ embedding dimension）H上执行的，H明显小于特征映射H/N × H/N的维度，其中N是2的因子，取决于块。在我们的实验中，除非另有说明，否则我们将H设为768。这种设计Tokenized MLP块的方法有助于编码有意义的特征信息，并且在计算或参数方面没有太多贡献。

训练过程

数据集：

International Skin Imaging Collaboration (ISIC 2018)

resize所有图像到512*512

Breast UltraSound Images (BUSI)

resize所有图像到256*256

实现细节：

损失函数：binary cross entropy (BCE) 和 dice loss

优化器(optimizer): Adam

学习率：0.0001，运用余弦退火学习率调度器，最小学习率可达0.00001

动量（momentum）：0.9

batch size： 8

训练轮次（epoch）：400

数据集划分比例：0.8：0.2

实验结果

总结

我们提出了一种新的深度网络架构UNeXt，用于医疗图像分割，专注于护理点应用。UNeXt是一种卷积和基于mlp的架构，其中有一个初始的conv阶段，然后是潜在空间中的mlp。具体来说，我们提出了一个带有移位MLP的标记化MLP块，以最小的复杂性和参数有效地建模表示。我们在多个数据集上验证了UNeXt，我们实现了更快的推理，降低了复杂性和更少的参数，同时还实现了最先进的性能。