RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition

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

论文：RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition

源代码：DingXiaoH/RepMLP

1 引入

本文提出了一个由多个全连接层构成的用于图像分类的模块RepMLP。全连接层能够高效地建模长距离依赖和位置模式，但不能很好地捕捉局部信息（擅长这件事情的是卷积）。我们在RepMLP中引入了卷积操作来捕捉局部信息，并在推理阶段将卷积核权重融入到全连接层中。

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

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

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

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

2 方法

在CNN中我们给卷积操作赋予了一个局部先验（也叫归纳偏置），卷积操作每次只处理相邻的元素。光有局部信息是不行的，需要给网络引入一定的全局信息。传统的CNN是通过卷积层的堆叠，来扩大感受野，以此获得全局信息。但这种做法效率低下，也带来很多优化的问题（以前训练很深的网络是很难的，后续通过正确的参数初始化和残差连接才逐步解决这一问题）。

此外CNN缺乏位置感知能力，因为一个卷积层在不同空间位置下是共享参数的。他没有考虑到各个空间位置特征的相对关系（Capsule Network也提到了）

最近的一些关于Vision Transformer的工作表明了在大量数据下，抛弃CNN的局部先验是可行的。本文尝试将全连接层替换部分卷积层中，以提供全局表征能力和位置感知能力。并将引入卷积层，赋予全连接层其不具备的捕捉局部信息能力。最后通过重参数化的方法，将卷积层和全连接层重参数化为一个全连接层，提升推理速度。

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

The input feature map is split into a set of partitions, and the Global Perceptron adds the correlations among partitions onto each partition. Then the Local Perceptron captures the local patterns with several conv layers, and the Partition Perceptron models the long-range dependencies. This sketch assumes $N = C = 1$,$H = W$, $H/w = W/w = 2$ (i.e., a channel is split into four partitions) for the better readability. We assume $ h, w > 7$ so that the Local Perceptron has conv branches of kernel size 1, 3, 5, 7. The shapes of parameter tensors are shown alongside FC and conv layers. Via structural re-parameterization, the training-time block with conv and BN layers is equivalently converted into a three-FC block, which is saved and used for inference.

2.1 Formulation

我们假设特征图表示为 ，我们采用F和W表示卷积与全连接层的核参数。为简单起见，本文采用了pytroch风格的数据排布与伪代码风格，比如 卷积处理的数据流表示如下： $$ M^{(o u t)}=C O N V\left(M^{(i n)}, F, p\right) $$

其中 表示输出特征，O表示输出通道数，p表示pad的像素数量， 表示卷积核(暂时假设groups=1)。为简单起见，我们还假设 。

对于全连接层，假设P和Q为输入与输出维度， 分别表示输入与输出， 表示全连接层核，计算公式如下： $$ V^{out} = MMUL(V^{in},W) = V^{in}\cdot W^T $$ 接下来我们聚焦于全连接层，假设FC不会改变输入输出分辨率，即 。我们采用RS(reshape)表示仅仅改变张量形变而不会改变内存数据顺序的操作，它是一种cost-free操作。输入首先会被平整为N个长度为 的向量，即 ，然后乘以核 得到输出 ，最后再RS为 。为更好的阅读，在歧义的情况下我们忽略RS而简写为： $$ M^{out} = MMUL(M^{in},W) $$

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

2.2 Global Perceptron

我们并未采用上述形式的全连接层，因为它不仅缺乏局部先验，而且会导致巨量参数： 。以通用配置为例， ，此时FC的参数量高达10G，这是完全不可接受的。为降低参数量，我们提出了Global Perceptron与Partition Peceptron分别进行分区间与分区内依赖性建模。

Global Perceptron 它将进行特征拆分且不同分区特征进行参数共享。比如， 将被拆分为 ，我们将每个 视作一个分区。假设 表示期望分区的高和宽，输入 首先reshape为 ，然后对坐标轴顺序进行调整 ，最后再reshape为 。通过这种方式，参数量可以从 下降到 。

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

• 对每一个切分后的特征图进行平均池化
• 送入BN层+两个全连接层
• 将向量reshape成$(\frac{NHW}{hw},C,1,1)$与切分后的特征图进行广播相加

相关代码为:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

def __init__(...): 
    if self.need_global_perceptron:
        internal_neurons = int(self.C * self.h_parts * self.w_parts // fc1_fc2_reduction)
        self.fc1_fc2 = nn.Sequential()
        self.fc1_fc2.add_module('fc1', nn.Linear(self.C * self.h_parts * self.w_parts, internal_neurons))
        self.fc1_fc2.add_module('relu', nn.ReLU())
        self.fc1_fc2.add_module('fc2', nn.Linear(internal_neurons, self.C * self.h_parts * self.w_parts))
        if deploy:
            self.avg = nn.AvgPool2d(kernel_size=(self.h, self.w))
        else:
            self.avg = nn.Sequential()
            self.avg.add_module('avg', nn.AvgPool2d(kernel_size=(self.h, self.w)))
            self.avg.add_module('bn', nn.BatchNorm2d(num_features=self.C))

def forward(...): 
    if self.need_global_perceptron:
        v = self.avg(inputs)
        v = v.reshape(-1, self.C * self.h_parts * self.w_parts)
        v = self.fc1_fc2(v)
        v = v.reshape(-1, self.C, self.h_parts, 1, self.w_parts, 1)
        inputs = inputs.reshape(-1, self.C, self.h_parts, self.h, self.w_parts, self.w)
        inputs = inputs + v

2.3 Partition Perceptron

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

Partition Perceptron 它包含FC与BN层，并以分区特征作为输入。前述输出 将通过reshape、re-arrange、reshape等操作变为 。我们进一步采用组卷及降低FC3的参数量，定义如下： $$ M^{(o u t)}=g C O N V\left(M^{(i n)}, F, g, p\right), F \in R^{O \times \frac{C}{g} \times K \times K} $$ 类似的，组形式的FC核 ，此时参数量可以减少g倍。尽管这种形式的FC不被某些框架(如Pytorch)支持，但我们可以采用 卷积代替实现。它的实现包含三个步骤：

• 将 reshape为空域尺寸为 的特征；
• 采用groups=g的 卷积；
• 将上述所得特征reshape为 。

整个过程定义如下：

相关代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

def __init__(): 
    self.fc3 = nn.Conv2d(self.C * self.h * self.w, self.O * self.h * self.w, 1, 1, 0, bias=deploy, groups=fc3_groups)
    self.fc3_bn = nn.Identity() if deploy else nn.BatchNorm1d(self.O * self.h * self.w)

def forward(): 
    #   Feed partition map into Partition Perceptron
    fc3_inputs = partitions.reshape(-1, self.C * self.h * self.w, 1, 1)
    fc3_out = self.fc3(fc3_inputs)
    fc3_out = fc3_out.reshape(-1, self.O * self.h * self.w)
    fc3_out = self.fc3_bn(fc3_out)
    fc3_out = fc3_out.reshape(-1, self.h_parts, self.w_parts, self.O, self.h, self.w)

2.4 Local Perceptron

Local Perceptron 它将分区特征经由几个卷积进行处理。前面的图示Fig1给出了 的示意图。理论上，仅有的约束为：核尺寸 (因为采用比分辨率更大的核没有意义)，但是我们参考ConvNet的常规操作仅采用了奇数核。为简单起见，我们采用 这种方框形式，其他 同样有效。所有卷积分支的输出与Partition perceptron的输出相加作为最终的输出。

需要注意的是为了后续的特征相加操作，这里卷积层都需要做padding。另外为了减少参数量，这里也是采用的分组卷积。这里的分组数需要与Partition Perceptron的分组数一致

具体做法如下：

• 将特征图同时送入到卷积核大小为1, 3, 5, 7的卷积层
• 送入BN层
• 与 Partition Perceptron 的输出相加

相关代码如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

def __init__(...): 
    self.reparam_conv_k = reparam_conv_k
    if not deploy and reparam_conv_k is not None:
        for k in reparam_conv_k:
            conv_branch = nn.Sequential()
            conv_branch.add_module('conv', nn.Conv2d(in_channels=self.C, out_channels=self.O, kernel_size=k, padding=k // 2, bias=False, groups=fc3_groups))
            conv_branch.add_module('bn', nn.BatchNorm2d(self.O))
            self.__setattr__('repconv{}'.format(k), conv_branch)

def forward(...): 
    #   Feed partition map into Local Perceptron
    if self.reparam_conv_k is not None and not self.deploy:
        conv_inputs = partitions.reshape(-1, self.C, self.h, self.w)
        conv_out = 0
        for k in self.reparam_conv_k:
            conv_branch = self.__getattr__('repconv{}'.format(k))
            conv_out += conv_branch(conv_inputs)
        conv_out = conv_out.reshape(-1, self.h_parts, self.w_parts, self.O, self.h, self.w)
        fc3_out += conv_out # fc3_out是前面Partition Perceptron的输出

2.5 Merging Conv into FC

在将RePMLP转换为三个FC之前，我们首先看一下如何将卷积合并到FC。假设FC核 ，卷积核 ，我们期望构建 满足： $$ M M U L\left(M^{(i n)}, W^{\prime}\right)=M M U L\left(M^{(i n)}, W^{(1)}\right)+C O N V\left(M^{(i n)}, F, p\right) $$ 我们注意到：对任意与 同形状的核 ，MMUL的加法特征满足： $$ M M U L\left(M^{(i n)}, W^{(1)}\right)+M M U L\left(M^{(i n)}, W^{(2)}\right)=M M U L\left(M^{(i n)}, W^{(1)}+W^{(2)}\right) $$ 因此，只要可以构建与 同形状的 ，我们就可以将F合并到 并满足： $$ M M U L\left(M^{(i n)}, W^{F, p}\right)=C O N V\left(M^{(i n)}, F, p\right) $$ 很明显， 一定存在(因为卷积可视作稀疏版FC)。考虑到不同平台对于卷积的加速策略、内存排布等方式的不同，A平台的矩阵构建可能并不适合于B平台。我们提出了一种简单的、平台无关解决方案。

正如前面所说，对于任意输入 ，卷积核F，存在一个FC核 满足： $$ M^{(o u t)}=\operatorname{CONV}\left(M^{(i n)}, F, p\right)=M M U L\left(M^{(i n)}, W^{(F, p)}\right) $$ 采用矩阵乘形式，此时有： $$ V^{(out)} = V^{(in)}\cdot W^{(F,p)_T} $$

我们对齐插入恒等矩阵 ，并采用如下定律： $$ V^{(out)} = V^{(in)}\cdot (I\cdot W^{(F,p)_T}) $$

注意： 是由F构建得到， 表示对特征 进行卷积。通过显示RS，此时有：

通过比较上述公式，我们可以得到： $$ W^{(F, p)}=R S\left(C O N V\left(M^{(I)}, F, p\right),(C h w, O h w)\right)^{T} $$ 该公式精确的展示了如何通过F，p构建 。简而言之，卷积核的等效FC核可以通过对恒等矩阵进行卷积并添加合适reshape得到。

2.6 Converting RepMLP into Three FC Layers

为采用上述理论，我们首先需要通过等价融合方式消除BN层。假设 表示卷积核， 表示BN中的参数。那么两者的融合方式如下：

3 实验结果

3.1 纯MLP和消融研究

The stage1 and stage3 are displayed in detail. Taking stage1 for example, $32×32$ is the resolution, $C = 16 $is the number of output channels (except the last layer). Left: FC(32,16) is the kernel size, suggesting that this FC (equivalent to a $1×1$ conv) projects 16 channels into 32 channels; all the RepMLPs are configured with $g = 2, h = w = 8$. Right: the convolutional counterpart uses $3×3$ conv. A BN follows every conv and a ReLU follows every RepMLP or conv-BN sequence.

在训练阶段，我们采用了标准数据增广(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实验的卷积，优于没有全局感知，精度仍出现了下降。

3.2 RepMLP-ResNet用于ImageNet分类

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

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

• 当仅仅对c4进行替换时，r=8时RepMLP-Res50具有比ResNet50更少的参数量，更快的推理速度(快10%)；
• Table2的前两行说明：当前深度学习框架对于组形式$1 \times 1$卷积支持程度并不够好，参数量提升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在速度于精度方面表现更优。

除了前述模块改进外，作者还提供了一种高速版本，其结果见上图右，性能见下表。

从上表可以看到：轻量版RepMLP取得了与原始ResNet50相当的进度，但FLOPs降低30%，推理速度快55%。

3.3 人脸识别

不同于卷积，FC不具备平移不变性性，这使得RepMLP非常适合于具有位置先验的图像(比如人脸)。因此，我们采用RepMLP在人脸识别任务上进行了有效性验证，结果见下表。

The speed (examples/second) is tested with a batch size of 512 and input 96×96 on the same 1080Ti GPU

从上表可以看到：

• 相比MobileFaceNet，FaceResNet具有更高的精度，但推理速度更慢；
• 相比MobileFaceNet，RepMLP-FaceRes取得了4.91%的精度提升，同时推理速度快8%(尽管具有更高的FLOPs)。

3.4 语义分割

语义分割是一种极具代表性的平移不变任务，我们在Cityscapes数据上验证RepMLP-Res50的有效性。结果见下表。

The speed (examples/second) is tested with a batch size of 16 and input 713×713 on the same 1080Ti GPU

从中可以看到：

• 相比Res-50，PSPNet+RepMLP-Res50-g4/8取得了2.31%的mIoU指标提升；
• 尽管RepMLP方案具有更多的参数量，但其FLOPs更低，推理速度更快。

参考资料

图解 RepMLP-技术圈

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