MLP-Mixer是ViT团队的另一个纯MLP架构的尝试。如果MLP-Mixer重新引领CV领域主流架构的话，那么CV领域主流架构的演变过程就是MLP->CNN->Transformer->MLP? 要回到最初的起点了吗？？？( Transformer移除了注意力以后就剩MLP了)
这篇论文提出了一种”纯“MLP结构的视觉架构。
先将输入图片拆分成patches，然后通过Per-patch Fully-connected将每个patch转换成feature embedding，然后送入N个Mixer Layer，最后通过Fully-connected进行分类。
Mixer分为channel-mixing MLP和token-mixing MLP两类。channel-mixing MLP允许不同通道之间进行交流；token-mixing MLP允许不同空间位置(tokens)进行交流。这两种类型的layer是交替堆叠的，方便支持两个输入维度的交流。每个MLP由两层fully-connected和一个GELU构成。
从上图我们可以看出，MLP -Mixer 首先使用图片分成很多个小正方形的patch,每个patch的大小定义为patch_size。论文中实现这一步骤使用的是前面提到的卷积，卷积核的大小和步长均patch_size。论文中给的参数，也是2的幂。
网络不再使用传统的RELU激活函数，而是使用了GELU激活函数。
将图片分成小块后，在将它转换为一维结构。如图：
然后将每一个patch进行转换，如下图所示：
通过这样一种方式呢，就将一张图片转换为了一个大矩阵，就可以输入到Mixer Layer 中进行计算。
MLP 是两个全连接层的感知机,W1,W2,对应token_mixer中两个全连接的权重，W3,W4则表示channel_mixer两个全连接的权重。σ表示GELU激活函数。那么公示就很简单了，输入X经过Layer Normalize,再乘以W1，再经过激活函数后乘以W2，再加上X。第二个公式也是相同的计算过程。
将前面通过编码得到的矩阵经过Layer Norm 在将矩阵进行旋转（T 表示旋转）连接MLP1,MLP1 就是文章token_mixer 用来寻找像素与像素之间的关系，其中，MLP1中的权值共享。计算完之后，再将矩阵旋转回来，通过Layer Norm 后再接一个channel_mixer 用于寻找通道与通道之间的关系。其中MixerLayer 还启用了ResNet中的跨连结构，跨连结构的作用可以参考[ResNet原理讲解和复现]，看到这里，是不是感觉它跟卷积的原理很类似。
从上图可以看出Mixer Layer的输入维度和输出维度相同，并且通过MLP的方式来寻找图片像素与像素，通道与通道的关系。
这就是MLP-MIXER的网络结构了
实现的难点在于，矩阵旋转，我们使用einops中的Rearrange实现矩阵旋转
使用Rearrange 实现旋转
Rearrange(‘b n d -> b d n’) #这里是[batch_size, num_patch, dim] -> [batch_size, dim, num_patch]
#定义多层感知机
import torch
import numpy as np
from torch import nn
from einops.layers.torch import Rearrange
from torchsummary import summary
import torch.nn.functional as F
class FeedForward(nn.Module):
def __init__(self,dim,hidden_dim,dropout=0.):
super().__init__()
self.net=nn.Sequential(
#由此可以看出 FeedForward 的输入和输出维度是一致的
nn.Linear(dim,hidden_dim),
#激活函数
nn.GELU(),
#防止过拟合
nn.Dropout(dropout),
#重复上述过程
nn.Linear(hidden_dim,dim),
nn.Dropout(dropout)
)
def forward(self,x):
x=self.net(x)
return x
class MixerBlock(nn.Module):
def __init__(self,dim,num_patch,token_dim,channel_dim,dropout=0.):
super().__init__()
self.token_mixer=nn.Sequential(
nn.LayerNorm(dim),
Rearrange('b n d -> b d n'), #这里是[batch_size, num_patch, dim] -> [batch_size, dim, num_patch]
FeedForward(num_patch,token_dim,dropout),
Rearrange('b d n -> b n d') #[batch_size, dim, num_patch] -> [batch_size, num_patch, dim]
)
self.channel_mixer=nn.Sequential(
nn.LayerNorm(dim),
FeedForward(dim,channel_dim,dropout)
)
def forward(self,x):
x=x+self.token_mixer(x)
x=x+self.channel_mixer(x)
return x
class MLPMixer(nn.Module):
def __init__(self,in_channels,dim,num_classes,patch_size,image_size,depth,token_dim,channel_dim,dropout=0.):
super().__init__()
assert image_size%patch_size==0
self.num_patches=(image_size//patch_size)**2
#embedding 操作，用卷积来分成一小块一小块的
self.to_embedding=nn.Sequential(nn.Conv2d(in_channels=in_channels,out_channels=dim,kernel_size=patch_size,stride=patch_size),
Rearrange('b c h w -> b (h w) c')
)
#经过Mixer Layer 的次数
self.mixer_blocks=nn.ModuleList([])
for _ in range(depth):
self.mixer_blocks.append(MixerBlock(dim,self.num_patches,token_dim,channel_dim,dropout))
self.layer_normal=nn.LayerNorm(dim)
self.mlp_head=nn.Sequential(
nn.Linear(dim,num_classes)
)
def forward(self,x):
x=self.to_embedding(x)
for mixer_block in self.mixer_blocks:
x=mixer_block(x)
x=self.layer_normal(x)
x=x.mean(dim=1)
x=self.mlp_head(x)
return x
MLP-Mixer用Mixer的MLP来替代ViT的Transformer，减少了特征提取的自由度，并且巧妙的可以交替进行patch间信息交流和patch内信息交流，从结果上来看，纯MLP貌似也是可行的，而且省去了Transformer复杂的结构，变的更加简洁，有点期待后续ViT和MLP-Mixer如何针锋相对的，感觉大组就是东挖一个西挖一个的，又把尘封多年的MLP给挖出来了
