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()

总结

欢迎留言交流讨论。

相关推荐
終不似少年遊*1 小时前
词向量与词嵌入
人工智能·深度学习·nlp·机器翻译·词嵌入
夏莉莉iy3 小时前
[MDM 2024]Spatial-Temporal Large Language Model for Traffic Prediction
人工智能·笔记·深度学习·机器学习·语言模型·自然语言处理·transformer
pchmi4 小时前
CNN常用卷积核
深度学习·神经网络·机器学习·cnn·c#
pzx_0014 小时前
【机器学习】K折交叉验证(K-Fold Cross-Validation)
人工智能·深度学习·算法·机器学习
造夢先森5 小时前
Transformer & LLaMA
深度学习·transformer·llama
神经美学_茂森6 小时前
神经网络防“失忆“秘籍:弹性权重固化如何让AI学会“温故知新“
人工智能·深度学习·神经网络
阿_旭6 小时前
【超详细】神经网络的可视化解释
人工智能·深度学习·神经网络
OpenCSG7 小时前
免部署|DeepSeek模型上线 CSGHub Serverless API|即刻免费试用
人工智能·深度学习·社区
CS_木成河8 小时前
【深度学习】预训练和微调概述
人工智能·深度学习·语言模型·微调·预训练