“重参数宇宙”再添新成员：RepMLP，清华大学&旷视科技提出将重参数卷积嵌入到全连接层

AIWalker

发布于 2021-05-24 10:42:38

7870

发布于 2021-05-24 10:42:38

文章被收录于专栏：AIWalkerAIWalker

标题&作者团队

paper: https://arxiv.org/abs/2105.01883

code: https://github.com/DingXiaoH/RepMLP（核心code已开源）

本文是清华大学&旷视科技在结构重参数领域继ACNet、RepVGG、DBB之后又一突破：RepMLP，它将重参数卷积与全连接层进行了巧妙地衔接，同时利用了全连接层的全局建模、位置感知特性与卷积的局部结构提取能力。在ImageNet分类、人脸识别以及语义分割等任务(无论是否具有平移不变性)上均能涨点。此外，相比标准卷积，RepMLP的一个优势：在大幅增加参数的同时不会造成推理速度的显著降低。比如，从RepMLP-Res50-g8/16到RepMLP-Res50-g4/8，参数量提升47%，但FLOPs仅提升3.6%，推理速度仅下降2.2%。

Abstract

本文提出一种多层感知器风格的神经网络构建模块RepMLP用于图像识别，它有一系列的全连接层构成。相比卷积层，全连接层更为高效，可以进行更好的长期依赖与位置模式建模，但在局部结构提取方面较差，因此通常不太适合于图像识别。

本文提出一种结构重参数技术，它为全连接层添加了局部先验信息以使其可以进行强有力的图像识别。具体来说：在训练阶段，我们在RepMLP内部构建了卷积层，而在推理阶段，我们将这些卷积层合并到全连接层内。

在CIFAR数据集上，简简单单的MLP模型即可取得与CNN非常接近的性能。通过将RepMLP插入到现有CNN中，我们在ImageNet数据集上提升ResNets模型精度达1.8%，在人脸识别任务上提升2.9%，在Cityscapes提升2.3%mIoU精度且具有更低的FLOPs。

该发现说明：全连接层的全局表达、位置感知能力与卷积的局部结构提取能力的组合能够以更快的速度在平移不变任务(比如语义分割)、图像对齐+位置模式相关任务(比如人脸识别)上的性能。

本文主要贡献包含以下几点：

提出利用全连接层的全局建模、位置感知能力，并为其插上局部先验的翅膀，使其适用于图像识别任务；
提出一种简单的、平台不可知的、可微分算法用于将并行卷积、BN合并到全连接层中，使其具有局部先验且不造成任何推理耗时增加；
提出一种高效模块RepMLP，并在多个视觉任务上证实了其有效性。

RepMLP

训练阶段的RepMLP（见下图）由三部分构成：Global Perceptron、Partition Perceptron、Local Perceptron。接下来，我们将针对所提RepMLP进行详细介绍，同时表明如何如何将训练阶段的RepMLP转换为推理阶段的全连接层。

Formulation

我们假设特征图表示为

M \in R^{N \times C \times H \times W}

，我们采用F和W表示卷积与全连接层的核参数。为简单起见，本文采用了pytroch风格的数据排布与伪代码风格，比如

K \times K

卷积处理的数据流表示如下：

M^{(out)} = CONV(M^{(in)}, F, p)

其中

M^{(out)} \in R^{N \times O \times H^{'} \times W^{'}}

表示输出特征，O表示输出通道数，p表示pad的像素数量，

F \in R^{O \times C \times K \times K}

表示卷积核(暂时假设groups=1)。为简单起见，我们还假设

H^{'} = H, W^{'} = W, stride=1, p=\lfloor \frac{K}{2}\rfloor

。

对于全连接层，假设P和Q为输入与输出维度，

V^{(in)} \in R^{N \times P}, V^{(out)} \in R^{N \times Q}

分别表示输入与输出，

W \in R^{Q \times P}

表示全连接层核，计算公式如下：

V^{(out)} = MMUL(V^{(in)}, W) = V^{(in)} \cdot W^T

接下来我们聚焦于全连接层，假设FC不会改变输入输出分辨率，即

H^{'} = H, W^{'} = W

。我们采用RS(reshape)表示仅仅改变张量形变而不会改变内存数据顺序的操作，它是一种cost-free操作。输入首先会被平整为N个长度为

CHW

的向量，即

V^{(in)} = RS(M^{(in)}, (N,CHW))

，然后乘以核

W(OHW, CHW)

得到输出

V^{(out)}(N,OHW)

，最后再RS为

M^{(out)(N.O,H,W)}

。为更好的阅读，在歧义的情况下我们忽略RS而简写为：

M^{(out)} = MMUL(M^{(in)}, W)

这种形式的全连接层无法充分利用图像的局部信息(因其忽视了位置信息)。

Components of RepMLP

我们并未采用上述形式的全连接层，因为它不仅缺乏局部先验，而且会导致巨量参数：

COH^2H^2

。以通用配置为例，

H=W=128,C=O=128

，此时FC的参数量高达10G，这是完全不可接受的。为降低参数量，我们提出了Global Perceptron与Partition Peceptron分别进行分区间与分区内依赖性建模。

Global Perceptron 它将进行特征拆分且不同分区特征进行参数共享。比如，

(N,C,14,14)

将被拆分为

(4N,C,7,7)

，我们将每个

7\times 7

视作一个分区。假设

h,w

表示期望分区的高和宽，输入

M\in R^{N\times C \times H \times W}

首先reshape为

(N,C,\frac{H}{h},h,\frac{W}{w},w)

，然后对坐标轴顺序进行调整

N, \frac{H}{h}, \frac{W}{w}, C, h, w

，最后再reshape为

(\frac{NHW}{hw}, C, h,w)

。通过这种方式，参数量可以从

COH^2W^2

下降到

COh^2w^2

。

然而，拆分将会打破相同通道不同分区之间的相关性。也就是说，每个分区将会进行独立建模，而忽视了位置上的相关性。为对每个分区添加相关性，Global Perceptron采用了以下几个操作：

对每个分区采用均值池化得到一个像素；
将上述所得送入到BN、两层MLP；
将上述所得进行reshape并与分区特征相加。

完成Global Perceptron操作后，将所得特征送入到后续的Partition Perceptron与Local Perceptron。

Partition Perceptron 它包含FC与BN层，并以分区特征作为输入。前述输出

(\frac{NHW}{hw}, C, h,w)

将通过reshape、re-arrange、reshape等操作变为

N,O,H,W

。我们进一步采用组卷及降低FC3的参数量，定义如下：

M^{(out)} = gCONV(M^{(in)}, F, g, p), F\in R^{O \times \frac{C}{g}\times K \times K}

类似的，组形式的FC核

W \in R^{Q\times \frac{P}{g}}

，此时参数量可以减少g倍。尽管这种形式的FC不被某些框架(如Pytorch)支持，但我们可以采用

1\times 1

卷积代替实现。它的实现包含三个步骤：

V^{(in)}

reshape为空域尺寸为

1\times 1

的特征；

采用groups=g的

1\times 1

卷积；

将上述所得特征reshape为

V^{(out)}

。

整个过程定义如下：

M^{'} = RS(V^{(in)}, (N, P, 1, 1)) \\ R^{'} = RS(W, (Q, \frac{P}{g}, 1, 1)) \\ gMMUL(V^{(in)}, W, g) = RS(gCONV(M^{'}, F^{'},g, 0), (N,Q))

Local Perceptron 它将分区特征经由几个卷积进行处理。前面的图示Fig1给出了

h,w>7, K=1,3,5,7

的示意图。理论上，仅有的约束为：核尺寸

K \le h, w

(因为采用比分辨率更大的核没有意义)，但是我们参考ConvNet的常规操作仅采用了奇数核。为简单起见，我们采用

K\times K

这种方框形式，其他

1\times3, 3\times5

同样有效。所有卷积分支的输出与Partition perceptron的输出相加作为最终的输出。

A Simple, Platform-agnostic, Differentiable Algorithm for Merging Conv into FC

在将RePMLP转换为三个FC之前，我们首先看一下如何将卷积合并到FC。假设FC核

W^{(i)} (Ohw, Chw)

，卷积核

F(O,C,K,K)(K \le h, w)

，我们期望构建

W^{'}

满足：

MMUL(M^{(in)}, W^{'}) = MMUL(M^{(in)}, W^{(1)}) + CONV(M^{(in)}, F, p)

我们注意到：对任意与

W^{(1)}

同形状的核

W^{(2)}

，MMUL的加法特征满足：

MMUL(M^{(in)}, W^{(1)}) + MMUL(M^{(in)}, W^{(2)}) = MMUL(M^{(in)}, W^{(1)} + W^{(2)})

因此，只要可以构建与

W^{(1)}

同形状的

W^{F,p}

，我们就可以将F合并到

W^{(1)}

并满足：

MMUL(M^{(in)}, W^{F,p}) = CONV(M^{(in)}, F, p)

很明显，

W^{(F,p)}

一定存在(因为卷积可视作稀疏版FC)。考虑到不同平台对于卷积的加速策略、内存排布等方式的不同，A平台的矩阵构建可能并不适合于B平台。我们提出了一种简单的、平台无关解决方案。

正如前面所说，对于任意输入

M^{(in)}

，卷积核F，存在一个FC核

W^{(F,p)}

满足：

M^{(out)} = CONV(M^{(in)}, F, p) = MMUL(M^{(in)}, W^{(F,p)})

采用矩阵乘形式，此时有：

V^{(out)} = V^{(in)} \cdot W^{(F,p)_T}

我们对齐插入恒等矩阵

I(Chw, Chw)

，并采用如下定律：

V^{(out)} = V^{(in)} \cdot (I \cdot W^{(F,p)_T})

注意：

W^{(F,p)}

是由F构建得到，

I\cdot W^{(F,p)_T}

表示对特征

M^{I}

进行卷积。通过显示RS，此时有：

M^{(I)} = RS(I, (Chw, C, h, w)) \\ I \cdot W^{(F, p)_T} = CONV(M^{(I)}, F, p) \\ V^{(out)} = V^{in} \cdot RS(I \cdot W^{(F, p)_T}, (Chw, Ohw))

通过比较上述公式，我们可以得到：

W^{(F, p)} = RS(CONV(M^{(I)}, F, p), (Chw, Ohw))^T

该公式精确的展示了如何通过F，p构建

W^{(F,p)}

。简而言之，卷积核的等效FC核可以通过对恒等矩阵进行卷积并添加合适reshape得到。

Converting RepMLP into Three FC Layers

为采用上述理论，我们首先需要通过等价融合方式消除BN层。假设

F \in R^{O \times \frac{C}{g}\times K \times K}

表示卷积核，

\mu, \sigma, \beta, \gamma

表示BN中的参数。那么两者的融合方式如下：

F^{'}_{i,:,:,:} = \frac{\gamma_i}{\sigma_i}F_{i,:,:,:} \\ b_I^{'} = - \frac{\mu_i \gamma_i}{\sigma_i} + \beta_i

1D形式BN于FC3的融合采用类似形式可得：

\hat{W} \in R^{Ohw \times \frac{Chw}{g}}, \hat{b} \in R^{Ohw}

。通过上述一系列的表换，我们就可以得到一个FC核于一个偏置向量，即为推理时的FC的参数。

Experiments

Pure MLP and Ablation Studies

我们首先在CIFAR10上验证RepMLP的有效性，测试的网络结构见下图。

在训练阶段，我们采用了标准数据增广(padding、RandomCrop、RandomFlip)，batch=128，cosine学习率衰减机制，训练了100epoch。结果见下表。

从上表可以看到：

纯MLP模型能够以52.8M FLPs达到91.11%的精度，不过该结果弱于Wide ConvNet；
转换前的计算量优于包含卷积与BN，其计算量非常大，这说明了结构重参数的重要性；
移除Local Perceptron，模型精度下降8.5%，说明了局部先验的重要性；
移除Gloabl Perceptron，模型精度下降1.5%，说明了全局建模的重要性；
替换FC3为卷积，尽管其感受野更大，但仍造成精度下降3.5.%，说明了FC要比卷积更有力；
采用RepMLP替换F实验的卷积，优于没有全局感知，精度仍出现了下降。

RepMLP-ResNet for ImageNet Classification

我们采用ResNet作为基线模型并用于评估RepMLP作为传统ConvNet构建模型的性能。模块定义见上图。

上面两个表给出了不同阶段替换时的性能对比，从中可以看到：

当仅仅对c4进行替换时，r=8时RepMLP-Res50具有比ResNet50更少的参数量，更快的推理速度(快10%)；
Table2的前两行说明：当前深度学习框架对于组形式

1\times1

卷积支持程度并不够好，参数量提升59%，但推理速度仅下降了0.7%；更进一步的组形式

1\times 1

优化有可能使得RepMLP更高效。

采用RepMLP模块替换ResNet中的模块会导致轻微的速度下降，但精度出现了显著提升。比如，仅仅采用RepMLP替换c4即可带来0.94%精度提升，参数量仅增加5M；c3+c4的替换可以取得最佳的均衡。

上表对比了更高输入分辨率下的性能对比。从中可以看到：

相比同参数量的传统ConvNet，RepMLP-Res50的计算量更低、推理速度更快。比如，相比

224\times 224

输入的ResNet101，RepMLP-Res50仅需50%FLOPs，更少的参数量，推理速度快50%，即可取得同等精度；当输入分辨率为

320\times 320

时，RepMLP-Res50在精度、速度以及FLOPs方面的优势更大。

提升RepMLP的参数量会导致非常轻微的速度下下降。比如，从RepMLP-Res50-g8/16到RepMLP-Res50-g4/8，参数量提升47%，但FLOPs仅提升3.6%，推理速度仅下降2.2%。这对于大型服务器的高吞吐量推理极为有用。
相比Nonlocal与GC，RepMLP-Res50的推理速度几乎相同，但精度高1%；
相比于GPU端低效的EfficientNet，RepMLP-Res50在速度于精度方面表现更优。