RepVGG 论文笔记

RepVGG: Making VGG-style ConvNets Great Again

RepVGG CVPR 2021

论文链接: https://arxiv.org/abs/2101.03697

一、 Problem Statement

目前的网络结构普遍有以下缺点:

• 复杂的分支结构: residual addtion, branch-concatenation等等，使得模型推理速度降低和减少内存的利用。
• 增加了MAC: depthwise convolution， channel shuffle都会增加内存访问消耗。

有很多因素影响推理速度，仅仅是参考FLOPs，不能够很准确反应实际的速度。

二、 Direction

VGG网络通过堆叠3x3卷积和ReLU的一个complain model，没有分支结构等等。

三、 Method

设计理念就是: Simple is Fast, Memory-economical, Flexible

• Fast
  目前很多多分支结构的网络有着比VGG-Net低的FLOPs值，但实际上它们运行起来并没有比VGG快。两个重要的因数: MAC(memory access cost)和并行化程度。 尽管branch addition或者concatenation的计算是忽略不计的，但是MAC就很重要了。MAC在groupwise convolution上占据了很大一部分的时间。另外一方面，并行化程度高的模型，在同等的FLOPs参数下，肯定是并行程度高的运行快。

• Memory-economical
  多分支模式是memory-inefficient的，因为所有的分支的数据都需要保留到addition或者concatenation之后，如下图所示:
  

  而plain模型只会在某一层使用数据，等到操作结束后，就会立即被释放。

• Flexible
  灵活性的问题。多分支模型限制了它的灵活性，因为每一个residual block的最后一个卷积层一定要和输入的大小相同。而且，多分值模型限制了channel prunning的使用，甚至会使得剪枝后导致性能下降等问题。

但多分支的网络模型能够避免梯度消失的问题，因此作者通过structural re-parameterizaton技术，使得训练的时候是多分支模型，推理的时候是plain模型。

基于上面设计理念，提出了下列方法:

1. Training-time Multi-branch Architecture

Plain模型有一个致命的缺点就是，性能差。受到了ResNet的启发，作者提出的re-parameterization技术。这个方法能够构建shortcut branch$y=x+f(x)$，当$x$的维度和residual block$f(x)$不匹配的时候，就变成了$y=g(x)+f(x)$，其中$g(x)$是convolutional shortcut，由1x1卷积实现的。

多分支结构的成功，如ResNet，是因为多分支架构使模型是众多较小模型的隐式集成。具体地说，对于$n$个blocks，模型可以解释为$2^n$个模型的集合，因为每个block将信息分流为两条路径。 但因为训练的时候，多分支模型是有好处的，而在推理的时候存在缺点，因此作者就只在训练的时候使用，具体结构如下:

$$y=x+g(x)+f(x)$$

效果如下:

2. Re-param for Plain Inference-time Model

那训练出来的网络怎么用于推理呢？假设使用$W^{(3)}\in {\mathbb{R}}^{C_2\times C_1\times 3\times 3}$表示3x3卷积核，$C_1$是输入通道数，$C_3$是输出通道数。$W^{(1)}\in {\mathbb{R}}^{C_2\times C1}$是1x1分支的卷积核。使用${\mu }^{(3)},{\sigma }^{(3)},{\gamma }^{(3)},{\beta }^{(3)}$作为3x3卷积层后面的BN层的学习参数。${\mu }^{(1)},{\sigma }^{(1)},{\gamma }^{(1)},{\beta }^{(1)}$作为1x1卷积层后面的BN层的学习参数。${\mu }^{(0)},{\sigma }^{(0)},{\gamma }^{(0)},{\beta }^{(0)}$是identity branch BN层的参数。让$M^{(1)}\in {\mathbb{R}}^{N\times C_1\times H_1\times W_1}$ 为输入，$M^{(2)}\in {\mathbb{R}}^{N\times C_2\times H_2\times W_2}$为输出。

如果$C_1=C_2,H_1=H_2,W_1=W_2$，则:
$$\begin{array}{rl}{M}^{(2)}& =bn(M^{(1)}\ast W^{(3)},{\mu }^{(3)},{\sigma }^{(3)},{\gamma }^{(3)},{\beta }^{(3)})\\ & +bn(M^{(1)}\ast W^{(1)},{\mu }^{(1)},{\sigma }^{(1)},{\gamma }^{(1)},{\beta }^{(1)})\\ & +bn(M^{(1)},{\mu }^{(0)},{\sigma }^{(0)},{\gamma }^{(0)},{\beta }^{(0)})\end{array}$$

第一步就是把BN和卷积进行融合
我们知道，卷积操作的公式为(忽略bias项):
$$Conv(M)=M\ast W$$
BN的公式为:
$$BN(M)=\gamma \ast \frac{(M-\mu )}{\sigma }+\beta$$
把BN和卷积融合成一个公式为:
$$BN(Conv(M))=\gamma \ast \frac{(M\ast W-\mu )}{\sigma }+\beta$$
所以进一步可以简化为:

$$BN(Conv(M))=\gamma \ast \frac{M\ast W}{\sigma }-\frac{\gamma \mu }{\sigma }+\beta$$

如果没有identity branch，即$\forall 1\le i\le C_2$，上面的式子就只有前面两项，且融合之后变为如下:
$$bn(M,\mu ,\sigma ,\gamma ,\beta {)}_{:,i,:,:}=(M_{:,i,:,:}-{\mu }_i)\frac{{\gamma }_i}{{\sigma }_i}+{\beta }_i$$

所以卷积的权重，和bias项可以转换为:
$$W_i^{{}^{\prime }}=\frac{{\gamma }_i}{{\sigma }_i}{W}_{i,:,:,:}{b}_i^{{}^{\prime }}=-\frac{{\mu }_i{\gamma }_i}{{\sigma }_i}+{\beta }_i$$

所以最终为:
$$bn(M\ast W,\mu ,\sigma ,\gamma ,\beta {)}_{:,i,:,:}=(M\ast W^{\prime }{)}_{:,i,:,:}+b_i^{{}^{\prime }}$$

第二步就是把卷积进行融合
我们知道3x3卷积核中心点，其实等于1x1卷积核。如果把1x1卷积核进行zero-padding，变成了3x3的卷积核，就可以进行融合了。那对于identity branch中，其实就是输入等于输出。也就是卷积核为1，分别对每个通道进行卷积，也就是depthwise卷积。 那就可以从1x1卷积核变为3x3。 所有的卷积核，都变成了3x3，那么就可以进行融合相加了，如下图右边所示。

3. 模型结构

看一下整体的网络结构:

作者不使用max pooling，而是全部使用3x3卷积。把整个网络框架设置成5个阶段，第一个阶段下采样，所以stride=2。

• 第一个阶段，数据输入分辨率大，因此只使用1个网络层。
• 最后一个阶段有较多的通道数，所以作者也只使用了一个网络层来节省参数。
• 把大部分网络层放在倒数第二个阶段。

设置了两个参数$a,b$，分别让他们设置为前四个阶段的通道数[64, 128,256, 512]和最后一个阶段的通道数。通常$b>a$。

四、 Conclusion

作者使用了re-parameterization技术，把训练和推理时候的机构进行区分，使得模型在ImageNet上达到了80%的Top-1 accuracy。

RepVGG 有以下优点:

• 没有任何分支，只是一个plain model。
• 和VGG一样，仅仅使用3x3卷积和ReLU。
• 没有复杂的结构设计，可以实现speed-accuracy的balance。

其实现的原理和TensorRT很类似，把Conv和BN融合等等。代码后续值得一看！

Reference

1. https://www.zhihu.com/zvideo/1380969462577876992