YOLOv5/v7 引入 RepVGG 重参数化模块

本篇博文代码出自YOLOv5-lite ，YOLOv5-lite的作者在CSDN的账号是 pogg_ ，大家可以关注一下，这也是一位在开源项目上做了很多工作的博主。

RepVGG的原理和融合推导过程可以看我的这篇博文：RepVGG：让VGG风格的ConvNets再次伟大
RepVGG思想融入到YOLOv5中的思想可以参考 pogg 的博文：Repvgg重参化对YOLO工业落地的实验和思考

重参数化算法原理

论文地址：https://arxiv.org/abs/2101.03697

我们提出了一种简单但功能强大的卷积神经网络结构，该模型在推理时类似于VGG，只有3×3的卷积和ReLU堆叠而成，而训练时间模型具有多分支拓扑结构。训练时间和推理时间结构的这种解耦是通过结构重新参数化技术实现的，因此该模型被命名为RepVGG。在ImageNet上，RepVGG达到了超过80%的TOP-1准确率，据我们所知，这是第一次使用普通模型。在NVIDIA 1080Ti GPU上，RepVGG型号的运行速度比ResNet-50快83%，比ResNet-101快101%，精度更高，并且与EfficientNet和RegNet等最先进的型号相比，显示出良好的精度和速度折衷。代码和经过训练的模型可在以下位置获得 https://github.com/megvii-model/RepVGG.

经典的卷积神经网络(ConvNet) VGG通过一个由conv、ReLU和pooling组成的简单体系结构在图像识别方面取得了巨大成功。随着Inception、ResNet和DenseNet的出现，大量的研究兴趣转移到了精心设计的架构上，使得模型越来越复杂，一些最近的架构是基于自动或手动架构搜索，或者搜索基于基本架构的混合尺寸策略得到的强大架构。

虽然许多复杂的卷积网络比简单的卷积网络具有更高的精度，但其缺点是明显的。

复杂的多分支设计（如ResNet中的残差相加和Inception中的分支连接）使模型难以实现和自定义，降低了推理速度和降低了内存利用率。
一些组件（例如Xception和MobileNets中的depth conv和ShuffleNets中的channel shuffle）增加了内存访问成本，缺乏各种设备的支持。

由于影响推理速度的因素太多，浮点运算（FLOPs）的数量并不能精确地反映实际速度。尽管一些新模型的FLOP低于老式模型，如VGG和ResNet-18/34/50，他们可能不会跑得更快，因此，VGG和ResNets的原始版本仍然大量用于学术界和工业界。
在本文中，我们提出了RepVGG，这是一种VGG风格的架构，其性能优于许多复杂的模型（图1）。RepVGG具有以下优点。

该模型具有类似VGG的无分支（即前馈）拓扑，这意味着每一层都将其唯一前一层的输出作为输入，并将输出馈送到其唯一后一层。
该模型的主体仅使用3×3 conv和ReLU。
具体的架构(包括特定的深度和层宽度)实例化时不需要自动搜索、手动细化、复合缩放，也不需要其他繁重的设计。

对于一个普通模型来说，要达到与多分支体系结构相当的性能水平是很有挑战性的。一种解释是，多分支拓扑，如ResNet，使模型成为众多浅层模型的隐式集成，从而训练多分支模型避免了梯度消失问题。

由于多分支结构的优点都是训练的，而缺点是不利于推理的，我们提出了通过结构重新参数化将训练时多分支结构和推理时平面结构解耦，即通过变换结构参数将结构从一个结构转换到另一个结构。具体地说，网络结构与一组参数相耦合，例如，

卷积层由四阶核张量表示。如果某一结构的参数可以转换为另一结构耦合的另一组参数，我们可以等效地用后者替代前者，从而改变整个网络架构。

具体来说，我们使用identity和1×1分支构造了训练时的RepVGG，这是受ResNet的启发，但采用了不同的方式，可以通过结构重新参数化来删除分支(图2、4)。经过训练后，我们用简单代数进行变换，将一个identity分支看作是一个降级的1×1 conv，后者可以进一步看作是一个降级的3×3 conv，这样我们就可以用原3×3 kernel、identity和1×1分支以及批归一化(BN)层的训练参数构造一个3×3 kernel。因此，转换后的模型有一堆3×3的 conv层，保存用于测试和部署。
在这里插入图片描述

在这里插入图片描述

参数量与计算量

模型	layers	parameters	gradients	GFLOPs
yolov5_RepVGG融合前	375	5574845	5574845	16.2
yolov5_RepVGG融合后	280	5390365	1180160	15.7

代码添加方式

第一步；将如下代码添加到common.py中：

# build repvgg block
# -----------------------------
def conv_bn(in_channels, out_channels, kernel_size, stride, padding, groups=1):
    result = nn.Sequential()
    result.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                        kernel_size=kernel_size, stride=stride, padding=padding, groups=groups,
                                        bias=False))
    result.add_module('bn', nn.BatchNorm2d(num_features=out_channels))

    return result


class SEBlock(nn.Module):

    def __init__(self, input_channels, internal_neurons):
        super(SEBlock, self).__init__()
        self.down = nn.Conv2d(in_channels=input_channels, out_channels=internal_neurons, kernel_size=1, stride=1,
                              bias=True)
        self.up = nn.Conv2d(in_channels=internal_neurons, out_channels=input_channels, kernel_size=1, stride=1,
                            bias=True)
        self.input_channels = input_channels

    def forward(self, inputs):
        x = F.avg_pool2d(inputs, kernel_size=inputs.size(3))
        x = self.down(x)
        x = F.relu(x)
        x = self.up(x)
        x = torch.sigmoid(x)
        x = x.view(-1, self.input_channels, 1, 1)
        return inputs * x


class RepVGGBlock(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size=3,
                 stride=1, padding=1, dilation=1, groups=1, padding_mode='zeros', deploy=False, use_se=False):
        super(RepVGGBlock, self).__init__()
        self.deploy = deploy
        self.groups = groups
        self.in_channels = in_channels

        padding_11 = padding - kernel_size // 2

        self.nonlinearity = nn.SiLU()

        # self.nonlinearity = nn.ReLU()

        if use_se:
            self.se = SEBlock(out_channels, internal_neurons=out_channels // 16)
        else:
            self.se = nn.Identity()

        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 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 forward(self, inputs):
        if hasattr(self, 'rbr_reparam'):
            return self.nonlinearity(self.se(self.rbr_reparam(inputs)))

        if self.rbr_identity is None:
            id_out = 0
        else:
            id_out = self.rbr_identity(inputs)

        return self.nonlinearity(self.se(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out))

    def fusevggforward(self, x):
        return self.nonlinearity(self.rbr_dense(x))
# repvgg block end
# -----------------------------

第二步；在yolo.py中找到 def fuse(self): ，用如下代码替换：

# --------------------------repvgg refuse---------------------------------
    def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
        print('Fusing layers... ')
        for m in self.model.modules():
            if type(m) is RepVGGBlock:
                if hasattr(m, 'rbr_1x1'):
                    kernel, bias = m.get_equivalent_kernel_bias()
                    rbr_reparam = nn.Conv2d(in_channels=m.rbr_dense.conv.in_channels,
                                                out_channels=m.rbr_dense.conv.out_channels,
                                                kernel_size=m.rbr_dense.conv.kernel_size,
                                                stride=m.rbr_dense.conv.stride,
                                                padding=m.rbr_dense.conv.padding, dilation=m.rbr_dense.conv.dilation,
                                                groups=m.rbr_dense.conv.groups, bias=True)
                    rbr_reparam.weight.data = kernel
                    rbr_reparam.bias.data = bias
                    for para in self.parameters():
                        para.detach_()
                    m.rbr_dense = rbr_reparam
                    m.__delattr__('rbr_1x1')
                    if hasattr(m, 'rbr_identity'):
                        m.__delattr__('rbr_identity')
                    if hasattr(m, 'id_tensor'):
                        m.__delattr__('id_tensor')
                    m.deploy = True
                    delattr(m, 'se')
                    m.forward = m.fusevggforward  # update forward
            if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
                m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
                delattr(m, 'bn')  # remove batchnorm
                m.forward = m.forward_fuse  # update forward
        self.info()
        return self
    # --------------------------end repvgg & shuffle refuse--------------------------------

第三步；在yolo.py中将RepVGGBlock添加到如下位置：

在这里插入图片描述
第四步；修改配置文件，yolov5_RepVGG配置文件如下：

# create by pogg
# parameters
nc: 80  # number of classes
depth_multiple: 1  # model depth multiple
width_multiple: 1  # layer channel multiple

# anchors
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5-repvgg backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [32, 3]],  # 0-P1/2
   [-1, 1, RepVGGBlock, [64, 3, 2]], # 1-P2/4
   [-1, 1, C3, [64]],
   [-1, 1, RepVGGBlock, [128, 3, 2]], # 3-P3/8
   [-1, 3, C3, [128]],
   [-1, 1, RepVGGBlock, [256, 3, 2]], # 5-P4/16
   [-1, 3, C3, [256]],
   [-1, 1, RepVGGBlock, [512, 3, 2]], # 7-P4/16
   [-1, 1, SPP, [512, [5, 9, 13]]],
   [-1, 1, C3, [512, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [128, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [128, False]],  # 13

   [-1, 1, Conv, [128, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [128, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [128, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [128, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [128, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [128, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]