【网络理解】RepVGG: Making VGG-style ConvNets Great Again

系列文章目录

前言

本文主要是对RepVGG的讲解
论文：https://arxiv.org/pdf/2101.03697.pdf
GitHub：https://github.com/DingXiaoH/RepVGG
参考博客

一、Netron模型可视化

git clone https://hub.fastgit.org/DingXiaoH/RepVGG.git

修改:test.py, 增加 convert to onnx。

def test():
    args = parser.parse_args()

    repvgg_build_func = get_RepVGG_func_by_name(args.arch)

    model = repvgg_build_func(deploy=args.mode=='deploy')

    model.eval()
    # convert to onnx
    print('=====convert to onnx=======')
    dummy_input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(model, dummy_input, 'RepVGG-A1.onnx', export_params=True, verbose=False, input_names=['input0'],
                      output_names=['output0'])
    print('=====convert sucess=======')
    exit()

按照README运行：

python test.py [imagenet-folder with train and val folders] train RepVGG-A1-train.pth -a RepVGG-BA1   # 这里选择A1

可以看到图中有4个None，意思是4个block不带identity，对应后面会有讲解

此时会生成训练结构RepVGG-A1.onnx，打开Netron，查看模型结构block截图如下图：

运行convert.py转换推理模型

python convert.py RepVGG-B2-train.pth RepVGG-A1-deploy.pth -a RepVGG-A1

再次运行test.py, 此时会生成推理结构RepVGG-A1-deploy.onnx，打开Netron，查看模型结构如下图：

python test.py [imagenet-folder with train and val folders] deploy RepVGG-A1-deploy.pth -a RepVGG-A1

二、代码实现

RepVGG有多个版本，通过num_blocks和width_multiplier来控制模型的大小，类似于yolov5的cfg。

接下来看 RepVGG 这个类， 我们看到由5个stage +一个自适应池化 + 全连接层组成
repvgg.py代码如下（示例）：

每个stage对应前面定义的num_blocks和width_multiplier

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)

通过print可以看到每个stage里block的组成

if __name__ == '__main__':
    RepVGG = create_RepVGG_A1()
    print(RepVGG)
    dummy_input = torch.randn(1, 3, 224, 224)
    output = RepVGG(dummy_input)

以stage2为例，如下可以看到不带 identity 对应第一幅图 Netron可视化的None，5个stage中一共四个block不带identity

(stage2): Sequential(
    (0): RepVGGBlock(
      (nonlinearity): ReLU()   # 不带 rbr_identity
      (rbr_dense): Sequential(
        (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
      (rbr_1x1): Sequential(
        (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): RepVGGBlock(
      (nonlinearity): ReLU()  # 带 rbr_identity
      (rbr_identity): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)  
      (rbr_dense): Sequential(
        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
      (rbr_1x1): Sequential(
        (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    

总结

简单看了一下模型结构，下次研究下RepVGG的算子融合策略