YOLOv5改进系列(十) 引入RepVGG重参数化模块

小酒馆燃着灯

已于 2024-01-29 01:13:12 修改

阅读量1k

点赞数 22

文章标签： YOLO 人工智能 pytorch 深度学习

于 2023-12-14 23:49:04 首次发布

原文链接：https://blog.csdn.net/weixin_43694096/article/details/128753663

版权

深度学习同时被 2 个专栏收录

213 篇文章 41 订阅

订阅专栏

YOLOv5改进

27 篇文章 29 订阅

订阅专栏

文章目录

重参数化算法原理
Batch Normalization
参数量与计算量
代码添加方式

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层，保存用于测试和部署。
在这里插入图片描述

在这里插入图片描述

Batch Normalization

在这里插入图片描述
我们发现经过激活函数之后，x的分布趋向于两侧，这样收敛的肯定越来越慢，所以我们想把他拉回正态分布。

参数量与计算量

模型	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):
	    # 检查分支是否为None，如果是，则返回0
	    if branch is None:
	        return 0, 0
	
	    # 如果分支是nn.Sequential类型
	    if isinstance(branch, nn.Sequential):
	        # 获取卷积核、滑动平均均值、滑动平均方差、gamma、beta和eps（批归一化中的参数）
	        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:
	        # 否则，分支是nn.BatchNorm2d类型
	        assert isinstance(branch, nn.BatchNorm2d)
	        if not hasattr(self, 'id_tensor'):
	            # 创建一个与输入通道数相匹配的3x3单位矩阵（identity 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）
	    std = (running_var + eps).sqrt()
	
	    # 计算t = gamma / std，将gamma扩展为与卷积核相同的形状
	    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'):
        # 如果存在rbr_reparam属性，应用非线性操作和SE模块
        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)

    # 返回非线性操作、SE模块和恒等映射的组合
    return self.nonlinearity(self.se(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out))

def fusevggforward(self, x):
    # 将输入通过非线性操作和rbr_dense分支
    return self.nonlinearity(self.rbr_dense(x))

# repvgg block end
# -----------------------------

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

# --------------------------repvgg refuse---------------------------------
def fuse(self):
    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_()
                
                # 替换原来的rbr_dense层为新的rbr_reparam层
                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
                
                # 删除SE模块
                delattr(m, 'se')
                
                # 更新模块的前向传播函数
                m.forward = m.fusevggforward
                
        if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
            # 融合Conv和BatchNorm层，更新Conv层
            m.conv = fuse_conv_and_bn(m.conv, m.bn)
            
            # 删除BatchNorm层
            delattr(m, 'bn')
            
            # 更新模块的前向传播函数
            m.forward = m.forward_fuse
    
    # 打印融合后的模型信息
    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)
  ]