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里怎么塞进RepVGG
YOLOv8的骨干是C2f模块,本质是跨阶段局部网络。我尝试了两种改造方案:
方案一:把C2f里的Bottleneck替换成RepVGGBlock。这个改动最小,但效果一般,因为RepVGGBlock本身没有跨层连接,和C2f的设计思路有点冲突。
方案二:直接替换整个C2f为RepVGGStage。这个改动大一些,但更彻底。具体做法是:把C2f的split操作去掉,用几个RepVGGBlock堆叠,每个Block之间用1×1卷积做通道变换。
我最终选了方案二,因为实测下来mAP高了0.7个点,参数量还少了15%。这里有个坑——RepVGGBlock的stride参数要小心处理。YOLOv8的骨干里有下采样层,如果直接把stride=2的RepVGGBlock放进去,合并后的卷积感受野会出问题。我的做法是单独用一个stride=2的3×3卷积做下采样,后面再接stride=1的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,不然梯度会传回去w=conv.weight mean=bn.running_mean.detach()var=bn.running_var.detach()gamma=bn.weight.detach()beta=bn.bias.detach()eps=bn.eps# 计算等效卷积权重和偏置std=(var+eps).sqrt()w_fused=w*(gamma/std).view(-1,1,1,1)b_fused=beta-gamma*mean/stdifconv.biasisnotNone:b_fused+=conv.bias*(gamma/std).view(-1)returnw_fused,b_fused这里有个细节:1×1卷积合并成3×3时,需要在四周补零。别这样写:
# 错误示范:直接reshapew_1x1=conv_1x1.weight.reshape(-1,1,3,3)# 这样不对!正确做法是用torch.nn.functional.pad:
w_1x1=torch.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骨干后:
- mAP@0.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,那个坑更多,到时候再聊。