QARepVGG--含demo实现

文章目录


前言

在上一篇博文RepVGG中,介绍了RepVGG网络。RepVGG 作为一种高效的重参数化网络,通过训练时的多分支结构(3x3卷积、1x1卷积、恒等映射)和推理时的单分支合并,在精度与速度间取得了优秀平衡。然而,其在低精度(如INT8)量化后常出现显著精度损失。

本文将要介绍的QARepVGG(Make RepVGG Greater Again: A Quantization-aware Approach)的提出正是为了解决这一问题。其核心贡献在于基础的Block设计:

引入

文章做了详细的消融实验来一步一步的推理出这种结构,本文在此不多做赘述。只大概提一下:RepVGG其实是由三个单元构成:权重、BN和ReLU。卷积操作一般不会影响权重值的改变,基本服从0~1分布;而根据BN层的公式,会出现一个乘法项,导致方差可能发生改变;另外,如果输入的数值范围很大,经过ReLU也会产生大的方差项,导致量化困难。

因此,QARepVGG去掉了BN层,并在三个分支后新加了一个BN层来将分布改成一个量化友好的分布。

当然,建议读者阅读原论文,好多实验的设计跟分析很透彻。

Demo实现

本文旨在复现一个QARepVGG Block,读者可一键运行:

python 复制代码
import torch
import torch.nn as nn

class QARepVGGBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        assert in_channels == out_channels, "输入输出通道必须相同!"
        
        # 分支1:3x3卷积 + BN
        self.conv3x3 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.bn3x3 = nn.BatchNorm2d(out_channels)
        
        # 分支2:1x1卷积(无BN)
        self.conv1x1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        
        # 分支3:恒等映射(无BN)
        self.identity = nn.Identity()  # 直接传递输入
        
        # 合并后的BN层
        self.final_bn = nn.BatchNorm2d(out_channels)
        
        # 初始化权重(关键!)
        self._init_weights()

    def _init_weights(self):
        """显式初始化权重"""
        nn.init.kaiming_normal_(self.conv3x3.weight, mode='fan_out', nonlinearity='relu')
        nn.init.zeros_(self.conv1x1.weight)  # 初始化为零,与恒等映射互补

    def forward(self, x):
        # 分支1:3x3卷积 + BN
        branch3x3 = self.bn3x3(self.conv3x3(x))
        # 分支2:1x1卷积
        branch1x1 = self.conv1x1(x)
        # 分支3:恒等映射
        branch_id = self.identity(x)
        # 合并后通过最终BN
        out = self.final_bn(branch3x3 + branch1x1 + branch_id)
        return out

    def reparameterize(self):
        """将多分支合并为单一3x3卷积,并融合BN参数"""
        # 1. 将各分支转换为等效3x3卷积
        # 分支1:3x3卷积 + BN3x3
        kernel3x3, bias3x3 = self._fuse_conv_bn(self.conv3x3, self.bn3x3)
        
        # 分支2:1x1卷积(无BN),填充为3x3
        kernel1x1 = self._pad_1x1_to_3x3(self.conv1x1.weight)
        bias1x1 = torch.zeros_like(bias3x3)  # 无偏置
        
        # 分支3:恒等映射(视为1x1单位矩阵卷积,填充为3x3)
        identity_kernel = torch.eye(self.conv3x3.in_channels, device=self.conv3x3.weight.device)
        identity_kernel = identity_kernel.view(self.conv3x3.in_channels, self.conv3x3.in_channels, 1, 1)
        kernel_id = self._pad_1x1_to_3x3(identity_kernel)
        bias_id = torch.zeros_like(bias3x3)
        
        # 2. 合并所有分支的权重和偏置
        merged_kernel = kernel3x3 + kernel1x1 + kernel_id
        merged_bias = bias3x3 + bias1x1 + bias_id
        
        # 3. 融合最终BN层参数
        scale = self.final_bn.weight / (self.final_bn.running_var + self.final_bn.eps).sqrt()
        merged_kernel = merged_kernel * scale.view(-1, 1, 1, 1)
        merged_bias = scale * (merged_bias - self.final_bn.running_mean) + self.final_bn.bias
        
        # 4. 构建合并后的卷积层
        merged_conv = nn.Conv2d(
            self.conv3x3.in_channels,
            self.conv3x3.out_channels,
            kernel_size=3,
            padding=1,
            bias=True
        )
        merged_conv.weight.data = merged_kernel
        merged_conv.bias.data = merged_bias
        return merged_conv

    def _fuse_conv_bn(self, conv, bn):
        """融合卷积和BN的权重与偏置"""
        kernel = conv.weight
        running_mean = bn.running_mean
        running_var = bn.running_var
        gamma = bn.weight
        beta = bn.bias
        eps = bn.eps

        std = (running_var + eps).sqrt()
        scale_factor = gamma / std

        fused_kernel = kernel * scale_factor.view(-1, 1, 1, 1)
        fused_bias = beta - running_mean * scale_factor
        return fused_kernel, fused_bias

    def _pad_1x1_to_3x3(self, kernel):
        """将1x1卷积核填充为3x3(中心为原权重,其余为0)"""
        if kernel.size(-1) == 1:
            padded = torch.zeros(kernel.size(0), kernel.size(1), 3, 3, device=kernel.device)
            padded[:, :, 1, 1] = kernel.squeeze()
            return padded
        return kernel 


def test_qarepvgg():
    torch.manual_seed(42)

    # 输入数据(小方差加速BN收敛)
    x = torch.randn(2, 3, 4, 4) * 0.1

    # 初始化模块
    block = QARepVGGBlock(3, 3)

    # 训练模式:更新BN统计量
    block.train()
    for _ in range(100):  # 充分训练
        y = block(x)
        y.sum().backward()  # 伪反向传播

    # 推理模式:合并权重
    block.eval()
    with torch.no_grad():
        # 原始输出
        orig_out = block(x)

        # 合并后的卷积
        merged_conv = block.reparameterize()
        merged_out = merged_conv(x)

    # 打印关键数据
    print("out:", orig_out.mean().item())
    print("merge:", merged_out.mean().item())
    print("diff:", torch.abs(orig_out - merged_out).max().item())

    # 验证一致性(容差1e-6)
    assert torch.allclose(orig_out, merged_out, atol=1e-6), f"合并失败!最大差值:{torch.abs(orig_out - merged_out).max().item()}"
    print("✅ 测试通过!")

test_qarepvgg()

总结

欢迎留言交流讨论。

相关推荐
hyshhhh14 小时前
【算法岗面试题】深度学习中如何防止过拟合?
网络·人工智能·深度学习·神经网络·算法·计算机视觉
Listennnn14 小时前
优雅的理解神经网络中的“分段线性单元”,解剖前向和反向传播
人工智能·深度学习·神经网络
牙牙要健康16 小时前
【目标检测】【深度学习】【Pytorch版本】YOLOV3模型算法详解
pytorch·深度学习·目标检测
Scc_hy17 小时前
强化学习_Paper_1988_Learning to predict by the methods of temporal differences
人工智能·深度学习·算法
誉鏐17 小时前
从零开始设计Transformer模型(1/2)——剥离RNN,保留Attention
人工智能·深度学习·transformer
神经星星18 小时前
无需预对齐即可消除批次效应,东京大学团队开发深度学习框架STAIG,揭示肿瘤微环境中的详细基因信息
人工智能·深度学习·机器学习
程序员Linc18 小时前
写给新人的深度学习扫盲贴:向量与矩阵
人工智能·深度学习·矩阵·向量
补三补四18 小时前
机器学习-聚类分析算法
人工智能·深度学习·算法·机器学习
荷包蛋蛋怪19 小时前
【北京化工大学】 神经网络与深度学习 实验6 MATAR图像分类
人工智能·深度学习·神经网络·opencv·机器学习·计算机视觉·分类
贤小二AI20 小时前
贤小二c#版Yolov5 yolov8 yolov10 yolov11自动标注工具 + 免python环境 GPU一键训练包
人工智能·深度学习·yolo