Mlp-Mixer简介及pytoch代码

Mlp-Mixer

论文: MLP-Mixer: An all-MLP Architecture for Vision.
pytorch 代码: An unofficial implementation of MLP-Mixer by pytorch.

Mlp-mixer是谷歌最近提出的基于纯mlp结构的cv框架，用来替代卷积和Transformer里的注意力操作。Mixer的架构如下所示：

网络结构

Mlp-Mixer包含如下三个部分：
1.嵌入层（per-patch linear embedding）：
   把HxW分辨率的图像分为S个不重叠的PxP的小块（patch），然后送入嵌入层，将小块的像素值用全连接映射到C维，则输出X是SxC的二维矩阵,S代表序列数（或token），C代表通道数。
2.混合层（MixerLayer）：
  混合层不会改变输入维度，即输出还是SxC。该层中包含token混合和通道混合两个全连接块（图右下），每个块中包含两层的全连接，中间采用GELU激活函数；除此之外，还采用了跨层直连、层归一化、dropout。
  混合层的结构为：
（1）先将SxC的输入对通道进行层归一化（对单个样本（token）的所有通道求均值和方差，然后归一）；
（2）其次经转置后CxS经token-mixing块，对相同通道的不同token进行空间上的混合（此处输入是对位置敏感的，所以不需要transformer的位置嵌入）；
（3）然后转置为SxC后与原始输入x进行跨层直连得到x2；
（4）最后经层归一化后，用channel-mixing对相同位置的不同通道进行混合，得到SxC的输出与x2跨层直连得到SxC的最终输出。
3.分类头（classfier head）：
  单层全连接，输入是SxC经全局平均池化（对各个channel求平均）后C维的向量，输出是num_classes维。

TIPS

1.per-patch linear embedding
  该层对不重叠的小块用全连接做嵌入，可用kernel_size和stride都为patch_size的卷积来替代，再把宽高展平和通道调换维度就得到SxC。
  全连接和卷积本质上可以相互转化，全连接的后一个神经元仅和前面临近的几个神经元稀疏连接，且共享权重时就成了卷积;
  卷积使用和图像大小相同的kernel_size且无stride时就变成了全连接。
2.channel-mixing&token-mixing
  类似深度可分离卷积，但是表达能力没有深度可分离卷积强。
  深度可分离卷积把传统卷积分成逐点卷积（pointwise）和逐通道卷积（depthwise），逐点卷积采用1x1的卷积核，改变了特征映射的通道数；然后逐通道卷积不同于传统卷积，对每个通道单独卷积，结果只改变图像分辨率不改变通道数。
  channel-mixing类似逐点卷积，对同一个token的不同通道进行混合，其实就是用了patch_size的卷积核（比较大），并在不同token之间共享权重;
  token-mixing类似逐通道卷积，对所有token的相同通道进行混合。但不同的是，token-mixing对不同通道都共享权重，而逐通道卷积不同通道的卷积核参数都不一样，所以说表达能力没有深度可分离卷积强。
3.GELU
  GELU激活函数即高斯误差单元（Gaussian Error Linear Unit），它是dropout、zoneout、relu属性的混合。其中relu对输入确定性的乘0或1；而dropout依概率对输出乘0；zoneout用于rnn中，依概率对上一个cell输出状态乘1，即随机失活这个cell，把上一个cell的状态向量跨cell直接输出到下一个cell。
  这些随机正则化器与非线性激活函数不同，他们不依赖于输入，并且辅助了非线性激活函数的效果。
  而GELU做的就是依赖于输入的随机正则化，即对x乘一个0或1的掩码mask，而这个mask是服从概率为Φ(x)的伯努利分布,即当随机变量X<=Φ(x)时mask=1，X>Φ(x)时mask=0，mask（m）为：

  而Φ(x)是正态分布的分布函数，严格递增。所以x越小，Φ(x)越小，则mask=1的概率就越小；反之x越大，mask=1的概率就越大。

  也就是说GELU依概率对x进行了缩放。具体来说，同层中某个神经元输出的x0比其他神经元的输出要大，则x0的缩放系数较大。
  GELU(x)可以被近似为：

  GELU,RELU,ELU（输出可为负的relu，加速训练）激活函数的图像如下所示：

总结

  总的来说，Mixer能够混合patch内部信息和不同patch的信息，不需要像Transformer那样结构复杂，且类似kernel_size比较大的卷积操作。期待之后各大佬能用Mixer在cv各个方面开花！

参考

[1]:MLP-Mixer: An all-MLP Architecture for Vision
[2]:GAUSSIAN ERROR LINEAR UNITS (GELUS)
[3]:https://www.jianshu.com/p/87bdeb6089e1
[4]:https://blog.csdn.net/liruihongbob/article/details/86510622
[5]:https://www.zhihu.com/question/332535296