RepVGG:训练多分支推理单路径的结构重参数化
FreeGuideOnline
最新
2026-06-21
什么是结构重参数化
结构重参数化是一种训练时使用复杂多分支结构,推理时等价转换为简单单路径模型的技术。它不改变模型的预测结果,但在训练阶段获得更强的学习能力,在部署阶段保持极致的推理速度。RepVGG 是该思想在卷积神经网络中的经典实现。
为什么需要训练与推理分离
- 多分支训练:残差连接、多尺度卷积能够缓解梯度消失,提供丰富的特征表达。
- 单路径推理:合并分支为单路卷积,避免分支间的同步开销,显著提升推理效率。
- 内存与速度:单路径结构的内存占用更少,硬件亲和性更高,在GPU、CPU和嵌入式设备上均有大幅加速。
RepVGG 的基本模块
RepVGG 的基础构建块在训练时由三个分支组成,推理时折叠成一个 3×3 卷积。
训练时的三分支结构
3×3卷积分支:标准卷积 + 批归一化(BN),提供主体特征提取能力。1×1卷积分支:卷积 + BN,增强非线性与跨通道信息交互。- 恒等分支(Identity):仅当输入输出通道数相同且空间尺寸一致时存在,实际上是一层不做处理的直连,等价于一个特殊的
1×1卷积核。
所有分支的输出相加后通过 ReLU 激活函数。
推理时的单路径折叠原理
利用卷积的线性和可加性,将三个分支的参数合并为一个等效的 3×3 卷积。
合并 1×1 卷积到 3×3
- 将
1×1卷积核视为中心权重非零、四周为零的3×3卷积核。 - 对
1×1卷积的 BN 参数做与卷积合并的处理(将 BN 融合进卷积权重和偏置)。 - 得到等价
3×3卷积的参数。
构造恒等分支的 3×3 卷积
- 恒等映射对应一个权重为单位矩阵的
1×1卷积。 - 为输入通道与输出通道一一对应,中心位置为 1,其余为 0。
- 同样将其视为
3×3卷积核(中心点权重为 1,周围为零)。 - 融合其对应的 BN(训练时的 Identity 分支也会跟随一个 BN 以避免初始训练不稳定)。
将三个 3×3 卷积相加
- 三个融合后的
3×3卷积核与偏置直接相加,得到最终的等效3×3卷积。 - 因为卷积运算满足加法分配律,最终输出与多分支完全一致。
- 最后接上 ReLU,形成单路径基本单元。
RepVGG 网络设计
整体架构
RepVGG 借鉴 VGG 的朴素设计哲学,仅使用 3×3 卷积和 ReLU,不引入 MaxPooling 以外的复杂操作。网络分为 5 个阶段,每个阶段由多个 RepVGG 块堆叠,通过 stride=2 的卷积进行下采样。
各阶段配置与缩放
RepVGG 提供不同深度和宽度的配置以适配多样算力,例如 RepVGG-A0、A1、A2、B0、B1、B2 等。缩放通过调节每层的通道数(宽度因子 a、b)和块数(深度因子)实现。
- 早期阶段通道数较少,参数量集中于深层。
- 第一层使用
stride=2的3×3卷积,后期通过池化或步幅下采样。 - 分类头采用全局平均池化加全连接层。
训练细节
- 批量归一化放置于卷积之后、加和之前,各分支独立 BN。
- 使用标准数据增强与余弦退火学习率调度。
- 重参数化仅发生在训练完成后、推理前,不参与梯度更新,因此训练时完全保留多分支梯度流。
重参数化的数学本质
核心是利用如下等价性:
对于 Conv(x) = W * x + b,当存在 BN 时,合并公式为:
W_fused = W * (gamma / sqrt(var + eps))
b_fused = (b - mean) * (gamma / sqrt(var + eps)) + beta
由于多分支形式为 sum_i (Conv_i(x)),则:
sum_i (Conv_i(x)) = (sum_i W_i_fused) * x + sum_i b_i_fused
因此将各分支的融合权重与偏置直接相加即可。1×1 和恒等映射都可先填充为 3×3 核,再套用上述规则。
部署转换的完整步骤
- 加载训练好的 RepVGG 多分支模型。
- 遍历每一个 RepVGG 块:
- 将
3×3分支的卷积和 BN 合并为等效3×3卷积。 - 将
1×1分支的卷积和 BN 合并,并将卷积核填充为3×3。 - 若存在恒等分支,构造中心为 1 的
3×3核并合并其 BN。 - 三者核与偏置相加,得到单路
3×3卷积。
- 将
- 删除所有多余的 BN 层和分支结构,替换为融合后的卷积。
- 保存为纯单路径模型。
- 验证前后输出误差应在数值精度范围内(通常小于 1e-5)。
实际效果与优势
- 精度:在同级别计算量下,RepVGG 达到甚至超越 ResNet、EfficientNet 等人工设计网络的准确率。
- 速度:GPU 推理时延较 ResNet 可降低 40% 以上,CPU 和专用推理芯片提升更为明显。
- 简易性:网络结构极度规整,仅包含一种算子和激活函数,易于调优和部署到各种框架。
- 与现代训练范式兼容:可与知识蒸馏、剪枝、量化等技术叠加使用。
代码实践要点
PyTorch 实现关键
- 训练模块定义
RepVGGBlock,包含rbr_3x3、rbr_1x1、rbr_identity三个子分支。 forward中执行out = self.rbr_3x3(x) + self.rbr_1x1(x) + (self.rbr_identity(x) if exists else 0),然后ReLU。- 部署转换函数
get_equivalent_kernel_bias:- 调用内部函数
_fuse_bn_tensor融合 3×3、1×1、identity 分支。 - 将 1×1 填充为 3×3:
F.pad(kernel, [1,1,1,1])。 - identity 的初始化权重需构造单位矩阵后同样填充。
- 最终求和。
- 调用内部函数
转换后模型构建
创建 RepVGGBackbone 类,构建时直接使用单一 nn.Conv2d 替换原 RepVGGBlock,并通过 load_state_dict 加载转换后的权重。
常见问题与注意事项
- 恒等分支仅在
in_channels == out_channels且stride=1时使用;否则该分支不参与。 - 如果输入输出通道不同,
1×1分支也具备调整通道的能力,恒等分支必须被省略。 - 转换后必须重新设置模型为
eval()模式。 - 部分框架(如 ONNX、TensorRT)需保证算子支持与融合后的卷积完全兼容。
- 对于分组卷积或深度可分离变体,重参数化规则需额外设计,不可直接套用。
扩展与前沿
- RepOptimizer:将结构重参数化与优化器重参数化结合,进一步分离训练和推理。
- DBB (Diverse Branch Block):引入多尺寸、多类型分支的更通用重参数化范式。
- RepNAS:利用神经架构搜索自动设计重参数化模块。
- 重参数化已从图像分类扩展到目标检测、分割和低层视觉任务,成为高效模型部署的重要工具。
结构重参数化通过巧妙的等价变换打破了多分支训练与快速推理之间的壁垒。RepVGG 以朴素而扎实的设计证明,在现代深度学习工程中,简单即高效。