深度学习入门|复现 Mlp-Mixer（pytorch and keras） tensorflow|深度学习|pytorch

文章目录

前言
一、mlp-mixer原理介绍
- 1.网络结构
- 2. MixerLayer
- 3.MLP-Mixer模型类型
二、网络实现
- 1.pytorch复现
- 2.keras 复现
总结

前言 mlp-mixer，Google又提出的一种基于感知机的网络。尽管CNN 已经在计算机视觉上取得很好的效果，最近提出来的基于Attention，以Vision Transformer 为首的神经网络已经 “杀疯” CV界，但是Google的大佬们认为CNN 和Attention 也不是必须的，于是就提出了mlp-mixer，在分类任务上也达到了很好的效果。但是网络的提出却早到了CNN之父LeCun的“教育”。因为网络的第一层（embedding时）却用到了卷积。

文章图片

LeCun认为，这不过是一个卷积核为1x1的卷积网络罢了。我在复现过程中，没有找到一个合适的数据集用该网络取得一个很好的效果。所以这里不附上训练的代码了。

项目	链接
论文	链接

博客中给出的就是网络结构的全部代码。
一、mlp-mixer原理介绍 1.网络结构

文章图片

从上图我们可以看出，mlp-mixer 首先使用图片分成很多个小正方形的patch,每个patch的大小定义为patch_size。论文中实现这一步骤使用的是前面提到的卷积，卷积核的大小和步长均patch_size。至于卷积核的输出通道数，就是自己拍一拍脑袋就决定了，例如512 768 1024等。
网络不再使用传统的RELU激活函数，而是使用了GELU激活函数。
将图片分成小块后，在将它转换为一维结构。如图：

文章图片

然后将每一个patch进行转换，如下图所示：

文章图片

通过这样一种方式呢，就将一张图片转换为了一个大矩阵，就可以输入到MixerLayer 中进行计算啦。
2. MixerLayer MixerLayer的结构如下图所示：

文章图片

我们看一下论文里给出的公示：

文章图片

MLP 是两个全连接层的感知机,W1,W2,对应token_mixer中两个全连接的权重，W3,W4则表示channel_mixer两个全连接的权重。σ表示GELU激活函数。那么公示就很简单了，输入X经过LN,再乘以W1，再经过激活函数后乘以W2，再加上X。第二个公式也是相同的计算过程。
大佬们都喜欢吧简单的问题复杂化吗？
将前面通过编码得到的矩阵经过Layer Norm 在将矩阵进行旋转（T 表示旋转）连接MLP1,MLP1 就是文章token_mixer 用来寻找像素与像素之间的关系，其中，MLP1中的权值共享。计算完之后，再将矩阵旋转回来，通过Layer Norm 后再接一个channel_mixer 用于寻找通道与通道之间的关系。其中MixerLayer 还启用了ResNet中的跨连结构，跨连结构的作用可以参考[优雅的复现ResNet50]，看到这里，是不是感觉它跟卷积的原理很类似。
从上图可以看出MixerLayer的输入维度和输出维度相同，并且通过MLP的方式来寻找图片像素与像素，通道与通道的关系。
这就是MLP-MIXER的网络结构了，目前的了解，没有开源的pytorch或者TensorFlow 预训练的权重。官方给出的代码和权重是基于JAX的。
谁找到了权重给我说一下嘤嘤嘤~~~
3.MLP-Mixer模型类型

文章图片

文章中给出的模型参数列表，Patch resolution 就是patch 的长宽。Hidden size 就是映射成前面提到的大矩阵的维度，Squence length 是计算后的结果。
以上图红色部分为例，输入图像大小为224*224，
然后分成的块大小为32*32.
那么（224*224）\(32*32)=7*7，Squence length 就为 49。Dc和Ds分别表示token_mixer和channel_mixer 中全连接层节点的个数。
这就是MLP-Mixer的全部过程了。
二、网络实现 1.pytorch复现实现的难点在于，矩阵旋转，我们使用einops中的Rearrange实现矩阵旋转。还需要使用torchsummary 来查看网络结构。安装：

pip install einops pip install torchsummary

首先我们来实现MLP 也就是FeedForward：

#定义多层感知机 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 #测试多层感知机 # mlp=FeedForward(10,20,0.4).to(device) # summary(mlp,input_size=(10,))

实现过程很简单，就是全连接结构
接着我们来实现Mixer Block，里面包含了 token_mixer 和channel_mixer，还有矩阵转置。

#使用Rearrange 实现旋转 Rearrange('b n d -> b d n') #这里是[batch_size, num_patch, dim] -> [batch_size, dim, num_patch]

实现如下：

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#测试mixerblock # x=torch.randn(1,196,512) # mixer_block=MixerBlock(512,196,32,32) # x=mixer_block(x) # print(x.shape)

更具上述定义好的网络零件，就可以实现我们最终的主网络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 #embedding 操作，看见没用卷积来分成一小块一小块的 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') ) 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 if __name__ == '__main__': 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))

depth=1，网络深度等于1 这样方便我们看整体结构，最后的运行结果：

文章图片

2.keras 复现 keras的复现过程与pytorch类似,但有几个注意的地方，网络中使用的GELU 激活函数在tensorflow>=2.4才可以使用，conda 国内镜像很难安装tensorflow 2.4 以上的GPU 版本。如果要使用的话tensorflow2.4一下版本的话，自定义GELU激活函数如下：

def gelu(x): cdf = 0.5 * (1.0 + tf.tanh( (np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3))))) return x * cdf

在keras中使用类的方法定义自己的网络层时，需要重写 get_config 函数不然模型无法保存。keras中借助Permute 层实现转置。
keras的完整实现如下：

from tensorflow import keras import tensorflow as tf import numpy as np from keras import backend as K from tensorflow.keras.layers import ( Add, Dense, Conv2D, GlobalAveragePooling1D, Flatten, Layer, LayerNormalization, Permute, Softmax, Activation, ) class MlpBlock(Layer): def __init__( self, dim: int, hidden_dim: int, activation=None, **kwargs ): super(MlpBlock, self).__init__(**kwargs)if activation is None: activation = keras.activations.geluself.dim = dim self.hidden_dim = dim self.dense1 = Dense(hidden_dim) self.activation = Activation(activation) self.dense2 = Dense(dim)def call(self, inputs): x = inputs x = self.dense1(x) x = self.activation(x) x = self.dense2(x) return xdef compute_output_shape(self, input_signature): return (input_signature[0], self.dim)def get_config(self): config = super(MlpBlock, self).get_config() config.update({'dim': self.dim, 'hidden_dim': self.hidden_dim }) return configclass MixerBlock(Layer): def __init__( self, num_patches: int, channel_dim: int, token_mixer_hidden_dim: int, channel_mixer_hidden_dim: int = None, activation=None, **kwargs ): super(MixerBlock, self).__init__(**kwargs)self.num_patches = num_patches self.channel_dim = channel_dim self.token_mixer_hidden_dim = token_mixer_hidden_dim self.channel_mixer_hidden_dim = channel_mixer_hidden_dim self.activation = activationif activation is None: self.activation = keras.activations.geluif channel_mixer_hidden_dim is None: channel_mixer_hidden_dim = token_mixer_hidden_dimself.norm1 = LayerNormalization(axis=1) self.permute1 = Permute((2, 1)) self.token_mixer = MlpBlock(num_patches, token_mixer_hidden_dim, name='token_mixer')self.permute2 = Permute((2, 1)) self.norm2 = LayerNormalization(axis=1) self.channel_mixer = MlpBlock(channel_dim, channel_mixer_hidden_dim, name='channel_mixer')self.skip_connection1 = Add() self.skip_connection2 = Add()def call(self, inputs): x = inputs skip_x = x x = self.norm1(x) x = self.permute1(x) x = self.token_mixer(x)x = self.permute2(x)x = self.skip_connection1([x, skip_x]) skip_x = xx = self.norm2(x) x = self.channel_mixer(x)x = self.skip_connection2([x, skip_x])# TODO need 2?return xdef compute_output_shape(self, input_shape): return input_shapedef get_config(self): config = super(MixerBlock, self).get_config() config.update({'num_patches': self.num_patches, 'channel_dim': self.channel_dim, 'token_mixer_hidden_dim': self.token_mixer_hidden_dim, 'channel_mixer_hidden_dim': self.channel_mixer_hidden_dim, 'activation': self.activation, }) return configdef MlpMixerModel( input_shape: int, num_classes: int, num_blocks: int, patch_size: int, hidden_dim: int, tokens_mlp_dim: int, channels_mlp_dim: int = None, use_softmax: bool = False, ): height, width, _ = input_shapeif channels_mlp_dim is None: channels_mlp_dim = tokens_mlp_dimnum_patches = (height*width)//(patch_size**2)# TODO verify how this behaves with same paddinginputs = keras.Input(input_shape) x = inputsx = Conv2D(hidden_dim, kernel_size=patch_size, strides=patch_size, padding='same', name='projector')(x)x = keras.layers.Reshape([-1, hidden_dim])(x)for _ in range(num_blocks): x = MixerBlock(num_patches=num_patches, channel_dim=hidden_dim, token_mixer_hidden_dim=tokens_mlp_dim, channel_mixer_hidden_dim=channels_mlp_dim)(x)x = Flatten()(x)# TODO verify this global average pool is correct choice herex = LayerNormalization(name='pre_head_layer_norm')(x) x = Dense(num_classes, name='head')(x)if use_softmax: x = Softmax()(x) return keras.Model(inputs, x)

keras 配置训练代码：
训练过程参数可以自行调试：

#学习率调试，首先我们设置一个较小的学习率查看loss的变化情况使用Tensorboard记录下来 import tensorflow as tf from keras_preprocessing.image import ImageDataGenerator from tensorflow.python.keras.callbacks import ModelCheckpointfrom tensorflow.keras.callbacks import (EarlyStopping, ReduceLROnPlateau, TensorBoard) from tensorflow.keras.optimizers import Adamdef train(): log_dir = './log' #训练日志路劲 train_dataset_path=r"your_tarin_data_path" #分类训练数据集路径 test_dataset_path=r"your_test_data_path" #分类测试集路径 batch_size = 64 # 加载数据集 lr= 1e-3 epochs=20 num_classes=1000 #你的分类数 train_datagen = ImageDataGenerator( #数据集增强，这些参数查阅keras 官方文档我前面的博客VGG 中说明过也有介绍说 rescale=1 / 255.0, rotation_range=20, zoom_range=0.05, width_shift_range=0.05, height_shift_range=0.05, shear_range=0.05, horizontal_flip=True, fill_mode="nearest", )train_generator = train_datagen.flow_from_directory( directory=train_dataset_path, target_size=(224, 224), color_mode="rgb", batch_size=batch_size, class_mode="categorical", shuffle=True, seed=42 ) test_datagen = ImageDataGenerator( rescale=1 / 255.0,)valid_generator = test_datagen.flow_from_directory(directory=test_dataset_path, target_size=(224, 224), color_mode="rgb", batch_size=batch_size, shuffle=True, seed=42 ) #你的模型，模型参数自己调试 mlp_mixer_base = MlpMixerModel(input_shape=(224, 224, 3), num_classes=num_classes, num_blocks=2, patch_size=16, hidden_dim=64, tokens_mlp_dim=32, channels_mlp_dim=64, use_softmax=True)mlp_mixer_base.summary() #training_weights='./weights'#这里是保存每次训练权重的如果需要自己取消注释 #checkpoint_period = ModelCheckpoint(training_weights + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5', #monitor='val_loss', save_weights_only=True, save_best_only=False, period=3) reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=4, verbose=1) #学习率衰减 early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1) # val_loss 不下降时停止训练防止过拟合 tensorboard = TensorBoard(log_dir=log_dir)#训练日志 optimizer=tf.keras.optimizers.Adam(learning_rate=lr) mlp_mixer_base.compile(loss=tf.keras.losses.categorical_crossentropy, metrics='acc',optimizer=optimizer) mlp_mixer_base.fit(train_generator,validation_data=https://www.it610.com/article/valid_generator, epochs=epochs,callbacks=[tensorboard, reduce_lr, early_stopping] ) mlp_mixer_base.evaluate(valid_generator,verbose=1) mlp_mixer_base.save('./mlp_mixer_base.h5') if __name__ == '__main__': train()

总结 【深度学习入门|复现 Mlp-Mixer（pytorch and keras）】代码复现过程，参考了论文地址给出的GitHub的链接，自己手撕代码的能力还是比较弱，不足之处就是没有使用数据集去训练它并达到一个不错的效果。后续如果有结果，会更新这篇博客，在训练模型时，调参过程是真的累，应该还是缺少理论知识的原因。后续会学习如何进行调参。
创作不易，点赞鼓励。
最后说一句，pytorch真香~~~~~~,前面的博客LeNet ，AlexNet和VGG已添加pytorch实现。