Mlp-mixer

• 论文：MLP-Mixer: An all-MLP Architecture for Vision

Mixer Architecture

MLP-Mixer主要包括三部分：Per-patch Fully-connected、Mixer Layer、分类器。

其中分类器部分采用传统的全局平均池化（GAP）+ 全连接层（FC）+ Softmax 的方式构成，故不进行更多介绍，下面主要针对前两部分进行解释。

Per-patch Fully-connected

FC相较于Conv，并不能获取局部区域间的信息，为了解决这个问题，MLP-Mixer通过Per-patch Fully-connected将输入图像转化为2D的Table，方便在后面进行局部区域间的信息融合。

具体来说，MLP-Mixer将输入图像相邻无重叠地划分为S个Patch，每个Patch通过MLP映射为一维特征向量，其中一维向量长度为C，最后将每个Patch得到的特征向量组合得到大小为S*C的2D Table。需要注意的时，每个Patch使用的映射矩阵相同，即使用的MLP参数相同。

实际上，Per-patch Fully-connected实现了（W，H，C）的向量空间到（S，C）的向量空间的映射。

例如，假设输入图像大小为2402403，模型选取的Patch为1616，那么一张图片可以划分为（240240）/（1616）= 225个Patch。结合图片的通道数，每个Patch包含了16163 = 768个值，把这768个值做Flatten作为MLP的输入，其中MLP的输出层神经元个数为128。这样，每个Patch就可以得到长度的128的特征向量，组合得到225128的Table。MLP-Mixer中Patch大小和MLP输出单元个数为超参数。

Mixer Layer

观察Per-patch Fully-connected得到的Table会发现，Table的行代表了同一空间位置在不同通道上的信息，列代表了不同空间位置在同一通道上的信息。换句话说，对Table的每一行进行操作可以实现通道域的信息融合，对Table的每一列进行操作可以实现空间域的信息融合。

在传统CNN中，可以通过1*1 Conv来实现通道域的信息融合，如果使用更大一点的卷积核，可以同时实现空间域和通道域的信息融合。

在Transformer中，通过Self-Attention实现空间域的信息融合，通过MLP同时实现空间域和通道域的信息融合。

而在MLP-Mixer中，通过Mixer Layer使用MLP先后对列、行进行映射，实现空间域和通道域的信息融合。与传统卷积不同的是，Mixer Layer将空间域和通道域分开操作，这种思想与Xception和MobileNet中的深度可分离卷积相似。

创新点

MLP-Mixer的创新点在于使用了MLP代替了Conv和Self-Attention，这表明使用最基础的MLP的拟合能力足以构建一个相对性能较好的CV模型。

模型图

代码实现

Mlp

import torch
import torch.nn as nn
from torchsummary import summary
from torch.nn import Conv2d
from einops.layers.torch import Rearrange, Reduce
from tensorboardX import SummaryWriter

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

Mixer-Layer

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'),
            FeedForward(num_patch,token_dim,dropout),
            Rearrange('b d n -> b n d')

         )
        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

Mlp-Mixer

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   # （224/16）**2=196
        # embedding 操作，看见没用卷积来分成一小块一小块的
        # 通过embedding可以将这张3*224*224的图片转换为Channel*Patches=512*196，再通过Rearrange转为196*512
        self.to_embedding=nn.Sequential(Conv2d(in_channels=in_channels,out_channels=dim,kernel_size=patch_size,stride=patch_size),
            Rearrange('b c h w -> b (h w) c')
        )

        # 输入为196*512的table
        # 以下为token-mixing MLPs（MLP1）和channel-mixing MLPs（MLP2）各一层
        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

测试

if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = MLPMixer(in_channels=3, dim=512, num_classes=1000, patch_size=16, image_size=224, depth=1, token_dim=256,
                     channel_dim=2048).to(device)
    summary(model,(3,224,224))

    # torch.Tensor([1, 2, 3, 4, 5, 6])
    inputs = torch.Tensor(1, 3, 224, 224)
    inputs = inputs.to(device)
    print(inputs.shape)

    # 将model保存为graph
    with SummaryWriter(log_dir='logs', comment='model') as w:
        w.add_graph(model, (inputs,))
        print("success")

结果

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1          [-1, 512, 14, 14]         393,728
         Rearrange-2             [-1, 196, 512]               0
         LayerNorm-3             [-1, 196, 512]           1,024
         Rearrange-4             [-1, 512, 196]               0
            Linear-5             [-1, 512, 256]          50,432
              GELU-6             [-1, 512, 256]               0
           Dropout-7             [-1, 512, 256]               0
            Linear-8             [-1, 512, 196]          50,372
           Dropout-9             [-1, 512, 196]               0
      FeedForward-10             [-1, 512, 196]               0
        Rearrange-11             [-1, 196, 512]               0
        LayerNorm-12             [-1, 196, 512]           1,024
           Linear-13            [-1, 196, 2048]       1,050,624
             GELU-14            [-1, 196, 2048]               0
          Dropout-15            [-1, 196, 2048]               0
           Linear-16             [-1, 196, 512]       1,049,088
          Dropout-17             [-1, 196, 512]               0
      FeedForward-18             [-1, 196, 512]               0
       MixerBlock-19             [-1, 196, 512]               0
        LayerNorm-20             [-1, 196, 512]           1,024
           Linear-21                 [-1, 1000]         513,000