RepVGG
简介
VGG网络是2014年牛津大学提出的, 在2014到2016年, VGG网络可以说是当时很火并广泛应用的backbone, 后面由于新网络的提出, 精度上VGG比不上ResNet, 速度和参数数量VGG比不过MobileNet等轻量级网络, 慢慢的VGG开始淡出人们的视线, 当VGG已经被大家遗忘时, 2021年清华,旷视等机构共同基础了RepVGG网络。
论文中, 作者提到了structural re-parameterization technique方法,即结构重参数化。实际上就是在训练时, 使用一个类似ResNet-style的多分支模型,而推理时转化成VGG-style的单路模型。 如下图, B表示RepVGG训练时采用的网络结构, 而在推理时采用图(C)的网络结构。
RepVGG Block详解
其实关于RepVGG模型就是在不断堆叠Rep VGG Block。下面介绍一下RepVGG Blocks中的结构, 如下图针对训练时采用的RepVGG Block结构。其中(a)是进行下采样stride=2时使用的RepVGG Block结构,图(b)是正常的(stride=1) RepVGG Block结构。通过图b可以发现, 训练时的RepVGG Block并行了三个分支, 一个卷积核大小为3x3的主分支, 一个卷积核大小为1x1的shortcu分支以及一个只连接了BN的shortcut分支。
为什么训练时采用多分支结构, 之前的Inception系列, ResNet以及DenseNet等模型, 可以发现这些模型都并行了多个分支, 根据现有的经验来看, 并行多个分支能够增加模型的表征能力。在论文中作者也简单做了消融实验, 在使用单路结构时,ACC大概为72.39, 在加上Identity branch以及1x1 branch后acc达到了75.14。
为什么推理时将多分支模型转换成单路模型, 论文中剃刀, 单路模型更快, 更省内存。
- 更快, 主要考虑到模型在推理时硬件的并行程度以及MAC(memory access cost), 对于多分支模型, 硬件需要分别计算每个分支的结果, 有的分支计算的快, 有的分支计算的慢, 而计算快的分支计算完后只能等其他分支计算完成后才能做进一步融合,这样会导致硬件算力不能充分利用, 或者说并行度不高。而且每个分支都需要方位一次内存, 计算完后还需要将计算结果存入内存(不断地访问和写入内存会在IO上浪费很多时间)
- 更省内存, 论文的图3中, 作者举了个例子, 如图A所以得Residual模块, 假设卷积层不改变channel的数量, 那么在主分支和shortcut分支上都要保存各自的特征图或者称Activation, 那么在add操作前占用的内存大概是输入activation的两倍,而图B的Plain结构占用内存始终不变。
- 更加灵活, 作者再论文中提到了模型优化的剪枝问题, 对于多分支的模型, 结构限制较多剪枝较麻烦, 而plain结构的模型就相对灵活很多, 剪枝也更方便。
其实除此之外, 在多分支转化成单路模型后很多算子进行了融合(比如conv2d和BN融合), 使得计算量变小了, 而且算子减少后启动kernel的次数也减少了(比如在GPU中, 每一次执行一个算子就要启动一次kernel, 启动kernel也需要消耗时间)。而且现在的硬件一般对3x3的卷积核做了大量的优化, 转成单路模型后采用的都是3x3卷积, 这样也能进一步加速推理。下图多分支模型B转换成单路模型C。
结构重参数化
在简单了解RepVGG Block的训练结构后, 下面看看RepVGG Block转成推理时的模型结构,即structural re-parametrization technique过程。根据论文中的图4可以看到, 结构重参数化主要分为两步, 第一步主要将Conv2d算子和BN算法融合以及将只有BN的分支转换成一个Conv2d算子, 第二步将每个分支上的3x3卷积层融合成一个卷积层。
融合Conv2d和BN
Conv2d和BN的融合对于网络的优化来讲已经是基本操作了。 因为conv2d和BN两个算子都是做线性运算, 所以可以融合成一个算子,这里需要强调一点, 融合是网络训练完之后做的, 所以现在讲的默认都是推理模型, 注意BN在训练以及推理时计算方式是不同的。对于卷积层, 每个卷积核的通道数与输入特征图的通道数相同, 卷积核的个数决定了输出特征图的通道个数。对于BN层(推理模式), 主要包含4个参数, μ(均值), σ2(方差), γ和β, 其中 μ(均值), σ2(方差)是训练过程统计得到的, γ和β是训练过程学习得到的。对于特征图第i个通道BN的计算公式如下, 其中为防止分母为0加上了一个非常小的数。
在论文3.3章节中, 作者给出了转换公式(对于通道i), 其中M代表输入BN层的特征图(activation), 这里忽略了上面分母加上的非常小的数。
所以转换后新的卷积层权重计算公式为(对于第i个卷积核), W‘和b’是新的权重和偏置。
为解释上面的过程, 作图如下:
假设输入的特征图(input feature map)如下图所示, 输入通道数为2, 然后采用两个卷积核(图中只画了第一个卷积核对应的参数)
计算输出特征图通道1上的第一个元素, 即当卷积核1在输入特征图红色框区域卷积时得到的额值(为保证输入输出特征图高宽不变, 所以对input feature map进行了padding), 其他位置的计算过程类似。
然后再将卷积层输出的特征图作为BN层的输入, 这里计算一下输出特征图通道1上的第一个元素, 按照上述BN在推理时的计算公式即可得到如下的计算结果。
将卷积层输出的特征图作为BN层的输入, 这里同样计算输出特征图通道1上的第一个元素。
最后对上述公式进行变形, 得到转化后新卷积层只需在对应第i个卷积核的权重上乘以
系数即可, 对应第i个卷积核新的偏置等于
因为之前采用Conv2d+BN的组合中Conv2d默认不采用偏置或偏置为0
Conv2d + BN融合实验
参考作者提供的源码, 首先了一个module包含了卷积核BN模块, 然后按照上述转换公式将卷积层的权重和BN的权重进行融合转换, 接着载入到新建的卷积模块fused_conv中, 嘴周随机创建一个Tensor(f1)将它们分别输入到module以及fused_conv中, 通过对比二者的输出可以发现它们的结果相同。
1 | # !/usr/bin/env python |
将1x1卷积换成3x3卷积
以1x1卷积层中的某一个卷积为例, 只需在原来权重周围补一圈零就行, 这样公式变成了3x3的卷积层, 为了保证输入输出特征图高宽不变, 此时需要将padding设置成1(原来卷积核大小为1x1时padding为0)
将BN换成3x3卷积
对于只有BN的分支由于没有卷积层, 所以我们可以先构建一个卷积层, 如下图, 构建一个3x3的卷积层, 该卷积层只做了恒等映射, 即输入输出特征图不变, 按照上述的融合方式将卷积层与BN层进行融合。
多分支融合
上面介绍了如何将每个分支融合转换成一个3x3的卷积层, 下面需进一步将多分支转换成一个单路3x3卷积层。
合并过程也很简单, 直接将这三个卷积层的参数相加即可。
接下来看论文的图就很清楚了。
模型配置
论文给中对模型进一步细分由RepVGG-A, RepVGG-B以及RepVGG-Bxgy三种配置。
可以看出RepVGG-B比RepVGG-A更深。RepVGG-A中的base layers of each stage为1, 2, 4, 14, 1, 而RepVGG-B为1, 4, 6, 16, 1。更加详细的配置可以看表3。其中a代表模型stage2-4的宽度缩放因子, b代表模型最后一个stage的宽度缩放因子。
而RepVGG-bxgy配置在RepVGG-B的基础上加入了组卷积(Group Convolution), 其中gy表示组卷积采用的groups参数为y, 注意不是所有的卷积层都采用组卷积, 根据源码可知, 从stage2开始(索引从1开始)的第2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26的卷积层采用组卷积。