MLP-Mixer的Pytorch实现和分析（一）

        MLP-Mixer网络结构分析：MLP-Mixer: An all-MLP Architecture for Vision_HHzdh的博客-CSDN博客

        Mixer的Pytorch代码实现的难点在于矩阵旋转，我们使用einops中的Rearrange实现矩阵旋转，同时使用torchsummary和tensorboardX来查看网络结构。

​

 1、导入相关库

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

 2、首先实现MLP（多层感知机），即上结构的MLP1和MLP2

        可以看到MLP的结构如下所示：Fully-connected、GELU、Fully-connected。 

​

# 这个是两层FC加一个激活函数的mlp block
# 因为有两个mixing,进出的维度都不变，只是中间全连接层的神经元数量不同
# 定义多层感知机
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
 

 2、接着来实现Mixer Block

​

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

 3、搭建MLPMixer

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

 4、测试MLPMixer

#测试Mlp-Mixer
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

        通过tensorboardX可以查看网络

 5、网络输入输出框图