021、RepVGG重参数化骨干：训练时多分支与推理时单路结构的YOLOv8实现

021、RepVGG重参数化骨干训练时多分支与推理时单路结构的YOLOv8实现从一次诡异的精度回退说起去年有个项目需要在边缘设备上跑YOLOv8n客户要求FPS不低于60。我把骨干换成ShuffleNetV2速度是上去了但mAP掉了将近4个点。换回原版C2f速度又不够。折腾了两周突然想起RepVGG那篇论文——训练时用多分支结构学得更充分推理时等价合并成单路既保精度又提速度。这个思路用在YOLOv8的骨干上简直是为边缘部署量身定做的。RepVGG的核心思想训练和推理是两套结构很多人第一次看RepVGG的代码会懵怎么forward里还有if判断其实它的设计哲学很简单训练时用3×3卷积1×1卷积恒等映射三条分支让梯度回传路径更丰富推理时把这三条分支等价合并成一个3×3卷积计算量直接砍掉三分之二。这里有个关键点——合并的前提是卷积核尺寸相同。1×1卷积和恒等映射怎么合并成3×3答案是补零。1×1卷积在3×3的kernel里只有中心点非零恒等映射等价于一个单位矩阵卷积同样可以填充成3×3的稀疏形式。BN层的合并更是常规操作把BN的缩放因子和偏置吸收进卷积权重里。YOLOv8里怎么塞进RepVGGYOLOv8的骨干是C2f模块本质是跨阶段局部网络。我尝试了两种改造方案方案一把C2f里的Bottleneck替换成RepVGGBlock。这个改动最小但效果一般因为RepVGGBlock本身没有跨层连接和C2f的设计思路有点冲突。方案二直接替换整个C2f为RepVGGStage。这个改动大一些但更彻底。具体做法是把C2f的split操作去掉用几个RepVGGBlock堆叠每个Block之间用1×1卷积做通道变换。我最终选了方案二因为实测下来mAP高了0.7个点参数量还少了15%。这里有个坑——RepVGGBlock的stride参数要小心处理。YOLOv8的骨干里有下采样层如果直接把stride2的RepVGGBlock放进去合并后的卷积感受野会出问题。我的做法是单独用一个stride2的3×3卷积做下采样后面再接stride1的RepVGGBlock。代码实现里的那些坑先看RepVGGBlock的核心实现。训练时forward是这样的defforward(self,x):ifself.training:# 训练时走多分支returnself.rbr_identity(x)self.rbr_1x1(x)self.rbr_3x3(x)else:# 推理时走合并后的单路returnself.rbr_reparam(x)注意这个self.training的判断PyTorch的model.train()和model.eval()会自动切换。但有个坑——如果你用torch.no_grad()做推理self.training还是True必须显式调用model.eval()。我之前在验证集上跑忘了切模式结果精度异常低排查了半天才发现是分支没合并。合并函数是重头戏deffuse_conv_bn(self,conv,bn):# 把BN的gamma和beta吸收进卷积权重# 这里踩过坑BN的running_mean和running_var要detach不然梯度会传回去wconv.weight meanbn.running_mean.detach()varbn.running_var.detach()gammabn.weight.detach()betabn.bias.detach()epsbn.eps# 计算等效卷积权重和偏置std(vareps).sqrt()w_fusedw*(gamma/std).view(-1,1,1,1)b_fusedbeta-gamma*mean/stdifconv.biasisnotNone:b_fusedconv.bias*(gamma/std).view(-1)returnw_fused,b_fused这里有个细节1×1卷积合并成3×3时需要在四周补零。别这样写# 错误示范直接reshapew_1x1conv_1x1.weight.reshape(-1,1,3,3)# 这样不对正确做法是用torch.nn.functional.padw_1x1torch.nn.functional.pad(conv_1x1.weight,[1,1,1,1])恒等映射更tricky。它等价于一个卷积核为单位矩阵的3×3卷积但输入输出通道必须相同。如果通道数不同这条分支直接去掉。实现时我踩过坑——直接用torch.eye生成单位矩阵但忘了考虑分组卷积的情况。YOLOv8里没有分组卷积所以还好。训练策略的调整换了RepVGG骨干后训练超参数需要微调。我发现几个关键点学习率要降低。原版YOLOv8用0.01的初始学习率RepVGG因为多分支结构梯度更丰富容易震荡。我降到0.005配合warmup稳定很多。权重衰减要加大。RepVGG的多分支结构天然有正则化效果但为了推理时合并后的权重更干净我把weight_decay从0.0005提到0.001。别加太多不然训练loss下不去。BN的momentum要调小。默认0.1对于RepVGG来说太大了因为多分支的BN统计量不稳定。我改成0.01让running_mean和running_var更新更平滑。推理时的合并时机合并操作在模型导出时做。我写了个工具函数遍历模型的所有模块遇到RepVGGBlock就调用merge方法。注意顺序——先合并BN再合并分支最后赋值给rbr_reparam。有个坑如果你用torch.jit.script或者onnx导出必须在导出前完成合并。因为script不支持动态的if self.training判断。我试过在forward里写条件分支结果torch.jit报错说无法解析条件表达式。合并后的模型可以直接用torch.save保存下次加载时就是单路结构。但如果你想保留训练能力建议保存两份权重——一份合并后的用于部署一份原始的多分支用于继续训练。实际效果在COCO数据集上用YOLOv8n做baseline替换RepVGG骨干后mAP0.5:0.95从37.3%涨到37.9%涨了0.6个点参数量从3.2M降到2.8M降了12.5%推理速度在TensorRT上从2.1ms降到1.7ms快了19%这个收益在轻量级模型上更明显。YOLOv8s从44.5%涨到45.2%参数量降了10%。但YOLOv8m以上收益递减因为大模型本身容量够大多分支带来的正则化效果不明显。个人经验RepVGG这个trick最适合的场景是边缘部署轻量级模型。如果你的模型已经很大了比如YOLOv8x换这个收益不大反而增加训练复杂度。训练时注意监控三个分支的梯度范数。如果某个分支的梯度一直很小说明这个分支没学到东西可以考虑去掉。我遇到过恒等映射分支梯度几乎为零的情况后来发现是通道数不匹配这条分支根本没起作用。合并后的模型对量化更友好。单路结构的权重分布更集中量化误差更小。我在INT8量化时RepVGG骨干的模型精度只掉了0.3个点原版C2f掉了0.8个点。最后说一句别在训练过程中做合并。有人想在每个epoch结束后合并一次再继续训练这会导致BN统计量混乱精度反而下降。训练和推理的结构必须严格分离。下一期准备写DenseNet风格的密集连接怎么融入YOLOv8的Neck那个坑更多到时候再聊。