040、RepVGG 重参数化 Block:3乘3加1乘1加Identity 训练时多分支到推理时单 3乘3 的等效合并

发布时间:2026/6/6 11:54:32

040、RepVGG 重参数化 Block:3乘3加1乘1加Identity 训练时多分支到推理时单 3乘3 的等效合并 040、RepVGG 重参数化 Block3乘3加1乘1加Identity 训练时多分支到推理时单 3乘3 的等效合并昨天调一个YOLOv8的改进实验发现改了Backbone之后推理速度反而比原版慢了30%。我盯着TensorRT的profiling结果看了半天发现瓶颈全在那些1×1卷积和残差连接的逐元素相加操作上——GPU对这类小算子并不友好频繁的kernel launch把延迟吃掉了。这让我想起RepVGG那篇论文里的一句话“多分支结构在训练时有用但在推理时是累赘。”问题本质训练和推理的矛盾做目标检测的都知道训练时我们希望模型有更强的表达能力多分支结构比如ResNet的残差连接、Inception的多尺度卷积能提供更好的梯度流动和特征融合。但推理时尤其是部署到边缘设备或使用TensorRT加速时我们希望模型结构越简单越好——最好全是3×3卷积加ReLU因为GPU和NPU对3×3卷积有深度优化。RepVGG的解决方案很暴力训练时用多分支推理前通过数学等价变换合并成单分支。这个合并过程就是重参数化Re-parameterization。训练时的Block结构先看训练时的RepVGG Block长什么样。一个标准的RepVGG Block包含三个并行分支classRepVGGBlock(nn.Module):def__init__(self,in_channels,out_channels,stride1):super().__init__()# 主分支3x3卷积 BNself.conv3x3nn.Conv2d(in_channels,out_channels,3,stride,padding1,biasFalse)self.bn3x3nn.BatchNorm2d(out_channels)# 辅助分支11x1卷积 BNself.conv1x1nn.Conv2d(in_channels,out_channels,1,stride,padding0,biasFalse)self.bn1x1nn.BatchNorm2d(out_channels)# 辅助分支2Identity恒等映射只有in_channelsout_channels且stride1时才有self.has_identity(in_channelsout_channelsandstride1)ifself.has_identity:self.bn_identitynn.BatchNorm2d(in_channels)defforward(self,x):# 三个分支分别计算后相加outself.bn3x3(self.conv3x3(x))outself.bn1x1(self.conv1x1(x))ifself.has_identity:outself.bn_identity(x)returnF.relu(out)这里有个细节容易踩坑Identity分支的BN层参数初始化。如果你直接复制上面的代码Identity分支的BN层默认的affine参数是随机初始化的这会导致训练初期梯度爆炸。正确做法是手动初始化BN的weight为1bias为0或者直接用nn.Identity()加一个单独的BN层——但这样又失去了重参数化的便利。我一般这样写# 别这样写nn.BatchNorm2d(in_channels) 默认参数可能不合适# 应该显式初始化ifself.has_identity:self.bn_identitynn.BatchNorm2d(in_channels)nn.init.constant_(self.bn_identity.weight,1.0)nn.init.constant_(self.bn_identity.bias,0.0)推理时的合并原理数学等价变换重参数化的核心思想卷积和BN可以合并成一个卷积多个卷积可以相加成一个卷积。第一步卷积BN合并训练时每个分支都是“卷积BN”的组合。推理时BN可以吸收进卷积的权重和偏置中。BN的推理公式是y gamma * (x - mean) / sqrt(var eps) beta卷积的输出是y W * x b合并后的卷积权重和偏置为W_merged gamma * W / sqrt(var eps) b_merged gamma * (b - mean) / sqrt(var eps) beta注意这里卷积的bias在训练时通常设为False因为后面有BN但合并后需要加上bias。第二步多分支卷积相加三个分支的卷积已经合并了BN可以相加成一个卷积。原理很简单卷积是线性操作多个卷积核在相同位置上的权重可以直接相加。但这里有个坑三个分支的卷积核尺寸不同。3×3卷积核是3×3的1×1卷积核是1×1的Identity分支相当于一个3×3的卷积核但中心为1、周围为0。所以合并时需要做padding操作1×1卷积核在周围补0变成3×3Identity分支构造一个3×3的卷积核中心为1其余为0代码实现合并函数这是我在YOLO改进中实际使用的合并函数注释里写满了踩过的坑deffuse_repvgg_block(block):将训练时的RepVGG Block合并成推理时的单3x3卷积# 第一步合并3x3分支的卷积和BN# 这里有个坑conv3x3的bias是False但合并后需要biaskernel_3x3,bias_3x3fuse_conv_bn(block.conv3x3,block.bn3x3)# 第二步合并1x1分支的卷积和BNkernel_1x1,bias_1x1fuse_conv_bn(block.conv1x1,block.bn1x1)# 1x1卷积核需要padding成3x3kernel_1x1pad_1x1_to_3x3(kernel_1x1)# 第三步处理Identity分支ifblock.has_identity:# Identity分支相当于一个3x3卷积中心为1周围为0kernel_id,bias_idfuse_identity_bn(block.bn_identity,block.in_channels)else:kernel_id,bias_id0,0# 第四步三个分支的权重和偏置相加fused_kernelkernel_3x3kernel_1x1kernel_id fused_biasbias_3x3bias_1x1bias_id# 构造合并后的单卷积层fused_convnn.Conv2d(block.in_channels,block.out_channels,3,strideblock.conv3x3.stride,paddingblock.conv3x3.padding,biasTrue)fused_conv.weight.datafused_kernel fused_conv.bias.datafused_biasreturnfused_conv其中fuse_conv_bn的实现要注意数值稳定性deffuse_conv_bn(conv,bn):# 别这样写直接计算忽略eps# 应该加上eps防止除零gammabn.weight.data betabn.bias.data meanbn.running_mean varbn.running_var epsbn.eps stdtorch.sqrt(vareps)# 这里踩过坑gamma和std的维度需要对齐gamma_stdgamma/std# 合并后的权重fused_weightconv.weight.data*gamma_std.view(-1,1,1,1)# 合并后的偏置fused_biasbeta-gamma*mean/stdreturnfused_weight,fused_bias在YOLO中的实际应用我在YOLOv8的Backbone中替换了部分C2f模块的Bottleneck为RepVGG Block。训练时保持多分支每个epoch结束后保存checkpoint时自动合并。推理时直接加载合并后的模型。这里有个性能数据在RTX 3060上原版YOLOv8s的Backbone推理延迟是2.3ms替换RepVGG后训练时延迟是3.1ms多分支开销但合并后降到1.9ms。合并后的3×3卷积比原版的1×13×3组合更快因为减少了kernel launch次数和内存访问。个人经验性建议不要在所有层都用RepVGG。我试过把整个Backbone都换成RepVGG Block训练时显存暴涨多分支需要保存更多中间变量而且梯度消失问题更严重。建议只在浅层前1/3的层使用深层保持原结构。Identity分支的添加条件要严格。只有输入输出通道相同且stride1时才加Identity否则会改变特征图尺寸。我见过有人强行在所有层加Identity结果stride2时特征图尺寸对不上训练直接崩。合并后的模型要验证精度。理论上合并是数学等价的但浮点运算的舍入误差会导致精度轻微下降通常在0.01%以内。我习惯在合并后跑一遍验证集确保mAP下降不超过0.1%。TensorRT部署时注意INT8量化。合并后的3×3卷积在INT8量化时可能比多分支更容易出现精度损失因为单层卷积的权重范围更大。建议先FP16部署确认没问题再尝试INT8。训练超参数需要微调。RepVGG的多分支结构相当于隐式增加了模型容量学习率可以适当降低比如从0.01降到0.008weight decay可以适当增加从5e-4增加到1e-3否则容易过拟合。最后说一句重参数化不是银弹。如果你的模型已经部署在CPU上多分支的逐元素相加反而可能比单卷积更快CPU对简单加法有优化。一定要先profiling再决定是否合并。

相关新闻