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 ) y=x+f(x) y=x+f(x)，当 x x x的维度和residual block f ( x ) f(x) f(x)不匹配的时候，就变成了 y = g ( x ) + f ( x ) y=g(x)+f(x) y=g(x)+f(x)，其中 g ( x ) g(x) g(x)是convolutional shortcut，由1x1卷积实现的。
多分支结构的成功，如ResNet，是因为多分支架构使模型是众多较小模型的隐式集成。具体地说，对于 n n n个blocks，模型可以解释为 2 n 2^n 2n个模型的集合，因为每个block将信息分流为两条路径。 但因为训练的时候，多分支模型是有好处的，而在推理的时候存在缺点，因此作者就只在训练的时候使用，具体结构如下:
y = x + g ( x ) + f ( x ) y = x + g(x) + f(x) y=x+g(x)+f(x)
效果如下:
2. Re-param for Plain Inference-time Model
那训练出来的网络怎么用于推理呢？假设使用 W ( 3 ) ∈ R C 2 × C 1 × 3 × 3 W^{(3)} \in \R^{C_2 \times C_1 \times 3 \times 3} W(3)∈RC2×C1×3×3表示3x3卷积核， C 1 C_1 C1是输入通道数， C 3 C_3 C3是输出通道数。 W ( 1 ) ∈ R C 2 × C 1 W^{(1)} \in \R^{C_2 \times C1} W(1)∈RC2×C1是1x1分支的卷积核。使用 μ ( 3 ) , σ ( 3 ) , γ ( 3 ) , β ( 3 ) \mu^{(3)}, \sigma^{(3)}, \gamma^{(3)}, \beta^{(3)} μ(3),σ(3),γ(3),β(3)作为3x3卷积层后面的BN层的学习参数。 μ ( 1 ) , σ ( 1 ) , γ ( 1 ) , β ( 1 ) \mu^{(1)}, \sigma^{(1)}, \gamma^{(1)}, \beta^{(1)} μ(1),σ(1),γ(1),β(1)作为1x1卷积层后面的BN层的学习参数。 μ ( 0 ) , σ ( 0 ) , γ ( 0 ) , β ( 0 ) \mu^{(0)}, \sigma^{(0)}, \gamma^{(0)}, \beta^{(0)} μ(0),σ(0),γ(0),β(0)是identity branch BN层的参数。让 M ( 1 ) ∈ R N × C 1 × H 1 × W 1 M^{(1)} \in \R^{N \times C_1 \times H_1 \times W_1} M(1)∈RN×C1×H1×W1 为输入， M ( 2 ) ∈ R N × C 2 × H 2 × W 2 M^{(2)} \in \R^{N \times C_2 \times H_2 \times W_2} M(2)∈RN×C2×H2×W2为输出。
如果 C 1 = C 2 , H 1 = H 2 , W 1 = W 2 C_1 = C_2, H_1=H_2, W_1=W_2 C1=C2,H1=H2,W1=W2，则:
M ( 2 ) = b n ( M ( 1 ) ∗ W ( 3 ) , μ ( 3 ) , σ ( 3 ) , γ ( 3 ) , β ( 3 ) ) + b n ( M ( 1 ) ∗ W ( 1 ) , μ ( 1 ) , σ ( 1 ) , γ ( 1 ) , β ( 1 ) ) + b n ( M ( 1 ) , μ ( 0 ) , σ ( 0 ) , γ ( 0 ) , β ( 0 ) ) \begin{aligned} M^{(2)} &= bn(M^{(1)}*W^{(3)}, \mu^{(3)}, \sigma^{(3)}, \gamma^{(3)}, \beta^{(3)}) \\ & +bn(M^{(1)}*W^{(1)}, \mu^{(1)}, \sigma^{(1)}, \gamma^{(1)}, \beta^{(1)}) \\ & + bn(M^{(1)}, \mu^{(0)}, \sigma^{(0)}, \gamma^{(0)}, \beta^{(0)}) \end{aligned} M(2)=bn(M(1)∗W(3),μ(3),σ(3),γ(3),β(3))+bn(M(1)∗W(1),μ(1),σ(1),γ(1),β(1))+bn(M(1),μ(0),σ(0),γ(0),β(0))
第一步就是把BN和卷积进行融合
我们知道，卷积操作的公式为(忽略bias项):
C o n v ( M ) = M ∗ W Conv(M) = M*W Conv(M)=M∗W
BN的公式为:
B N ( M ) = γ ∗ ( M − μ ) σ + β BN(M) = \gamma * \frac{(M-\mu)}{\sigma} + \beta BN(M)=γ∗σ(M−μ)+β
把BN和卷积融合成一个公式为:
B N ( C o n v ( M ) ) = γ ∗ ( M ∗ W − μ ) σ + β BN(Conv(M)) = \gamma * \frac{(M*W-\mu)}{\sigma} + \beta BN(Conv(M))=γ∗σ(M∗W−μ)+β
所以进一步可以简化为:
B N ( C o n v ( M ) ) = γ ∗ M ∗ W σ − γ μ σ + β BN(Conv(M)) = \gamma * \frac{M*W}{\sigma}- \frac{\gamma \mu}{\sigma} + \beta BN(Conv(M))=γ∗σM∗W−σγμ+β
如果没有identity branch，即 ∀ 1 ≤ i ≤ C 2 \forall 1 \leq i \leq C_2 ∀1≤i≤C2，上面的式子就只有前面两项，且融合之后变为如下:
b n ( M , μ , σ , γ , β ) : , i , : , : = ( M : , i , : , : − μ i ) γ i σ i + β i bn(M, \mu, \sigma, \gamma, \beta)_{:,i,:,:} = (M_{:,i,:,:} - \mu_i) \frac{\gamma_i}{\sigma_i} + \beta_i bn(M,μ,σ,γ,β):,i,:,:=(M:,i,:,:−μi)σiγi+βi
所以卷积的权重，和bias项可以转换为:
W i ′ = γ i σ i W i , : , : , : b i ′ = − μ i γ i σ i + β i W^{'}_i=\frac{\gamma_i}{\sigma_i}W_{i,:,:,:} \quad b^{'}_i = -\frac{\mu_i \gamma_i}{\sigma_i} + \beta_i Wi′=σiγiWi,:,:,:bi′=−σiμiγi+βi
所以最终为:
b n ( M ∗ W , μ , σ , γ , β ) : , i , : , : = ( M ∗ W ′ ) : , i , : , : + b i ′ bn(M * W, \mu,\sigma,\gamma,\beta)_{:,i,:,:} = (M * W')_{:,i,:,:} + b^{'}_i bn(M∗W,μ,σ,γ,β):,i,:,:=(M∗W′):,i,:,:+bi′
第二步就是把卷积进行融合
我们知道3x3卷积核中心点，其实等于1x1卷积核。如果把1x1卷积核进行zero-padding，变成了3x3的卷积核，就可以进行融合了。那对于identity branch中，其实就是输入等于输出。也就是卷积核为1，分别对每个通道进行卷积，也就是depthwise卷积。 那就可以从1x1卷积核变为3x3。 所有的卷积核，都变成了3x3，那么就可以进行融合相加了，如下图右边所示。
3. 模型结构
看一下整体的网络结构:
作者不使用max pooling，而是全部使用3x3卷积。把整个网络框架设置成5个阶段，第一个阶段下采样，所以stride=2。
- 第一个阶段，数据输入分辨率大，因此只使用1个网络层。
- 最后一个阶段有较多的通道数，所以作者也只使用了一个网络层来节省参数。
- 把大部分网络层放在倒数第二个阶段。
设置了两个参数 a , b a, b a,b，分别让他们设置为前四个阶段的通道数[64, 128,256, 512]和最后一个阶段的通道数。通常 b > a b \gt a b>a。
四、 Conclusion
作者使用了re-parameterization技术，把训练和推理时候的机构进行区分，使得模型在ImageNet上达到了80%的Top-1 accuracy。
RepVGG 有以下优点:
- 没有任何分支，只是一个plain model。
- 和VGG一样，仅仅使用3x3卷积和ReLU。
- 没有复杂的结构设计，可以实现speed-accuracy的balance。
其实现的原理和TensorRT很类似，把Conv和BN融合等等。代码后续值得一看！
Reference
- https://www.zhihu.com/zvideo/1380969462577876992