Vision MLP –ResMLP

Feedforward networks for image classification with data-efficient training

我们提出了ResMLP，一种完全基于多层感知机(MLP)进行图像分类的体系结构。 它是一个简单的残差网络，它交替（i）线性层，其中图像 patches在通道之间独立且相同地交互；以及（ii）两层前馈网络，其中通道中的每个 patch独立地相互作用。

CODE:

import torch
import numpy as np
from resmlp import ResMLP

img = torch.ones([1, 3, 224, 224])

model = ResMLP(in_channels=3, image_size=224, patch_size=16, num_classes=1000,
                 dim=384, depth=12, mlp_dim=384*4)

parameters = filter(lambda p: p.requires_grad, model.parameters())
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
print('Trainable Parameters: %.3fM' % parameters)

out_img = model(img)

print("Shape of out :", out_img.shape)  # [B, in_channels, image_size, image_size]

本文作者提出了一种基于全连接层的图像分类网络。网络结构与MLP-Mixer相似，即先将输入图像拆分成若干patch，对每个patch通过全连接层转换为特征嵌入矩阵，该矩阵的两个维度分别表示channel维度(每个局部位置的特征维度)和patch维度(表示局部位置的维度)。首先将该矩阵转置后沿patch维度进行全连接层运算，实现不同patch之间的交互；再沿channel维度进行全连接运算，实现不同channel之间的交互。最后使用池化层和输出层获得分类结果。本文与MLP-Mixer的不同之处在于采用了更强的数据增强方法和蒸馏策略。

当采用现代的训练策略进行训练时，使用大量的数据增广和可选的蒸馏方法，可以在ImageNet上获得令人惊讶的良好精度/复杂度折衷。

Affine仿射变换：

函数名称：diag(x)
函数功能：构建一个n维的方阵，它的主对角线元素值取自向量x，其余元素都为0