RepVGG是一个VGG风格的分类模型，使用了3x3的卷积和ReLu作为基本单元，并且在训练和推理时使用了不同的结构。作者在论文中提到该模型在ImageNet中达到了80%的top-1 acc，并且比ResNet50、ResNet101、EfficientNet等模型更快。作者称之为Make VGG Greate Again并用在了论文题目上，川宝听了直呼内行。
offical github：https://github.com/megvii-model/RepVGG
offical github：https://github.com/DingXiaoH/RepVGG
paper：RepVGG: Making VGG-style ConvNets Great Again
原理
一般来说，制约深度学习模型推理性能的因素主要为以下两点：
-
floating-point operations，即FLOPs。代表模型做一次推理所需要计算的浮点数。 -
memory access cost，即MAC。代表模型在推理时访问DRAM时带来的损耗。
VGG，ResNet之类的经典模型属于高flops，低mac；而MobileNet之类的模型属于低flops，高mac。尤其是在group-wise conv一类的操作中，mac带来的性能损耗占了相当大的比重。这也是为什么MobileNet、ShuffleNet之类的模型有着较小的理论FLOPS和参数量，但是其推理性能并没有等比提升的原因。
回到VGG和ResNet，ResNet在VGG上的主要改进就是将直上直下的单分支模型改成了带有residual block的多分支模型。这样做的好处主要是解决了模型在训练过程中，梯度传播随着模型深度的加深而消散的问题，使我们能够使用较小的参数量训练出较深的并且可收敛的模型。
本文的作者认为，residual block的好处主要体现在训练过程中，而在推理过程中residual block则带来了两个影响推理速度的因素：一是深度模型每多一层或一个分支，模型就要增加一次内存访问，相应的增大了mac；二是多分支的模型结构，主分支需要等到所有的子分支都计算完成才能继续进行下一步的推理，这造成了模型在推理过程中总会有一定的等待时间。因此RepVGG的主要思路就是参考ResNet构建一个多分支的模型，不同的是多分支的部分使用1x1的卷积跳层实现，这个结构主要在训练过程中使用。在推理过程中使用Re-parameterization的方法将block分支整合到主分支的卷积层中以此得到一个单分支的模型。
参考上图，RepVgg的一个block结构可以用下面的公式进行表示。其中M1代表block的intput，M2代表block的output，W3代表3x3的卷积核，W1代表1x1的卷积核，bn代表BatchNorm，μ、σ、γ、β代表BN层的滑动均值、标准差、缩放因子和偏置。可以看出RepVgg的一个完整block结构由3个分支构成。
在进行计算图优化的时候，我们通常会进行mergeBN的操作来削减BN层的操作，论文里也用到了这个方法进行分支的合并，下面稍微做一下相关的推导。
通常卷积、BatchNorm和激活函数是放在一起作为一个卷积单元来使用的。我们知道多个线性运算的层叠加在一起还是线性运算，本质上可以用一个线性运算的层来进行表示，因此卷积层和BatchNorm层可以合并为一个新的卷积层，从而减少一个卷积单元的计算量。
卷积的本质是一个线性运算，其计算公式如下所示：
BatchNorm的计算也是一个加减乘除的线性运算，其计算公式如下所示：
可以简化为：
卷积层的输出y作为BatchNorm的输入x带入公式，可以得到一个新的w和b：
新的卷积层可以表达为：
论文中用下面的公式对BN和Conv进行了合并的操作，本质上和我们上面的推导是一致的：
通过上面的变换，可以得到一个3x3的卷积核，两个1x1的卷积核和3个偏置。然后将两个1x1的卷积核padding到3x3和3x3的卷积核相加，将3个偏置相加，就得到了一个3x3卷积层+BN层所需的所有参数了。需要注意的是channel和stride必须对齐，如下图所示：
模型
作者通过控制宽度和深度获得了不同大小的RepVgg模型，如下：
下面是RepVGG的基础结构RepVGGBlock，其中rbr_identity是1x1的 BN identity，rbr_dense是一个3x3的卷积+BN，rbr_1x1是一个3x3的卷积+BN，在训练阶段输出三者的和，在推理阶段则数据合并后的卷积，通过deploy参数进行控制。
class RepVGGBlock(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size,
stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros', deploy=False):
super(RepVGGBlock, self).__init__()
self.deploy = deploy
self.groups = groups
self.in_channels = in_channels
assert kernel_size == 3
assert padding == 1
padding_11 = padding - kernel_size // 2
self.nonlinearity = nn.ReLU()
if deploy:
self.rbr_reparam = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
padding=padding, dilation=dilation, groups=groups, bias=True, padding_mode=padding_mode)
else:
self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels and stride == 1 else None
self.rbr_dense = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups)
self.rbr_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding_11, groups=groups)
print('RepVGG Block, identity = ', self.rbr_identity)
def forward(self, inputs):
if hasattr(self, 'rbr_reparam'):
return self.nonlinearity(self.rbr_reparam(inputs))
if self.rbr_identity is None:
id_out = 0
else:
id_out = self.rbr_identity(inputs)
return self.nonlinearity(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out)
下面是推理模式下合并参数的操作。首先通过_fuse_bn_tensor函数对三个分支都进行merge BN的操作得到weight和bias，注意里面也考虑了组卷积的情况。然后将只有BN层的rbr_identity分支通过_pad_1x1_to_3x3_tensor函数padding到3x3。最后三者相加得到最终的weight和bias。这里1x1的卷积没做padding是因为已经在上面的卷积操作中做过了。
# This func derives the equivalent kernel and bias in a DIFFERENTIABLE way.
# You can get the equivalent kernel and bias at any time and do whatever you want,
# for example, apply some penalties or constraints during training, just like you do to the other models.
# May be useful for quantization or pruning.
def get_equivalent_kernel_bias(self):
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
def _pad_1x1_to_3x3_tensor(self, kernel1x1):
if kernel1x1 is None:
return 0
else:
return torch.nn.functional.pad(kernel1x1, [1,1,1,1])
def _fuse_bn_tensor(self, branch):
if branch is None:
return 0, 0
if isinstance(branch, nn.Sequential):
kernel = branch.conv.weight
running_mean = branch.bn.running_mean
running_var = branch.bn.running_var
gamma = branch.bn.weight
beta = branch.bn.bias
eps = branch.bn.eps
else:
assert isinstance(branch, nn.BatchNorm2d)
if not hasattr(self, 'id_tensor'):
input_dim = self.in_channels // self.groups
kernel_value = np.zeros((self.in_channels, input_dim, 3, 3), dtype=np.float32)
for i in range(self.in_channels):
kernel_value[i, i % input_dim, 1, 1] = 1
self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
kernel = self.id_tensor
running_mean = branch.running_mean
running_var = branch.running_var
gamma = branch.weight
beta = branch.bias
eps = branch.eps
std = (running_var + eps).sqrt()
t = (gamma / std).reshape(-1, 1, 1, 1)
return kernel * t, beta - running_mean * gamma / std
def repvgg_convert(self):
kernel, bias = self.get_equivalent_kernel_bias()
return kernel.detach().cpu().numpy(), bias.detach().cpu().numpy(),
下面是模型的主体结构，完全由RepVGGBlock构成。在推理阶段，不考虑全连接层的话可以认为整个模型只有3x3的卷积这一种结构。其中num_blocks用来控制每个channel对应的block堆叠的数量，也就是控制整个模型的深度；width_multiplier用来控制每个block的channel数，也就是控制整个模型的宽度。
class RepVGG(nn.Module):
def __init__(self, num_blocks, num_classes=1000, width_multiplier=None, override_groups_map=None, deploy=False):
super(RepVGG, self).__init__()
assert len(width_multiplier) == 4
self.deploy = deploy
self.override_groups_map = override_groups_map or dict()
assert 0 not in self.override_groups_map
self.in_planes = min(64, int(64 * width_multiplier[0]))
self.stage0 = RepVGGBlock(in_channels=3, out_channels=self.in_planes, kernel_size=3, stride=2, padding=1, deploy=self.deploy)
self.cur_layer_idx = 1
self.stage1 = self._make_stage(int(64 * width_multiplier[0]), num_blocks[0], stride=2)
self.stage2 = self._make_stage(int(128 * width_multiplier[1]), num_blocks[1], stride=2)
self.stage3 = self._make_stage(int(256 * width_multiplier[2]), num_blocks[2], stride=2)
self.stage4 = self._make_stage(int(512 * width_multiplier[3]), num_blocks[3], stride=2)
self.gap = nn.AdaptiveAvgPool2d(output_size=1)
self.linear = nn.Linear(int(512 * width_multiplier[3]), num_classes)
def _make_stage(self, planes, num_blocks, stride):
strides = [stride] + [1]*(num_blocks-1)
blocks = []
for stride in strides:
cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
blocks.append(RepVGGBlock(in_channels=self.in_planes, out_channels=planes, kernel_size=3,
stride=stride, padding=1, groups=cur_groups, deploy=self.deploy))
self.in_planes = planes
self.cur_layer_idx += 1
return nn.Sequential(*blocks)
def forward(self, x):
out = self.stage0(x)
out = self.stage1(out)
out = self.stage2(out)
out = self.stage3(out)
out = self.stage4(out)
out = self.gap(out)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
实验
根据作者的实验，在acc接近的情况下， RepVgg在FLOPS、FPS和winograd乘法操作量等指标上都有显著的提升。
为了测试其真实的推理性能，我在NVIDIA RTX 3070+CUDA11+Pytorch1.8的环境性进行了测试。根据论文选用模型RepVGG-B1g2和resnet50 进行对比，分类数为1000，输入尺寸为(224， 224， 3)。假设Acc和论文中提到的一样，只有1%左右的差距，实际测试RepVGG相比ResNet50的推理性能基本上提升了一倍。并且由于其全3x3卷积的结构特性，其TensorRT的加速效果应该会优于resnet50。关于模型的Acc测试后面有时间我会训练一下看看。
| type | ResNet50 | RepVGG-B1g2 |
|---|---|---|
| average | 7.26ms/138FPS | 3.59ms/278FPS |
| fastest | 5.98ms/176FPS | 2.96ms/337FPS |
| slowest | 15.92ms/63FPS | 8.97ms/111FPS |