资讯详情

资讯详情

RepVGG结构重参数化:图像分类模型训练推理解耦与加速实战

简介压缩包内含2000个文件共986.61MB其中以2435个png图片为数据主体配合12个Python源码、2个pth权重文件、2个json配置及1个txt说明构成一套完整的RepVgg图像分类实战项目。项目围绕“VGG式”结构展开即无分支的plain架构、仅用3×3卷积和ReLU激活适合希望深入理解重参数化模型并落地图像分类任务的PyTorch开发者。已有990人学习下载。压缩包内附带可复现的训练代码、预训练权重、数据集和配置文件按目录组织清晰便于对照官方博客逐步实践。通过运行脚本可直接体验从数据预处理到模型训练、验证与推理的完整流程也可基于现有权重快速测试效果对掌握RepVgg的设计思路和工程实现有直接帮助。1. RepVgg在图像分类里到底解决了什么问题做过图像分类落地的工程师基本都撞过同一堵墙既要精度又要推理速度。模型改大一点FPS就掉得没法看模型压小精度又跟着崩。RepVgg就是为这个矛盾给出的一个工业级解法——它用“训练时多分支、推理时单路”的结构重参数化思路让网络在训练阶段拥有近似多分支集成的表达能力推理时再无损融合成一条纯3x3卷积链把精度和推理开销都占了。即便放在最新的图像分类模型对比里RepVgg的推理时延也常常比同精度区间的经典模型低一到两倍而且代码可读性比NAS搜出来的结构友好得多。这套方案特别适合接下来这类场景边缘设备部署、实时分类服务、以及对算子兼容性敏感的生产环境。接下来的内容会把理论、训练、转换、避坑一次讲透目标是你照着复现能跑通换到自己的数据集上也不翻车。2. 训练与推理的解耦结构重参数化的原理和网络搭建2.1 为什么训练要多分支推理却要单路径传统CNN比如ResNet解决梯度退化靠的是残差连接推理时残差结构还在相当于每层都要做一次额外的跳跃相加。而RepVgg的思路很直接训练阶段让每个Block内部同时走三条路径——一个3x3卷积、一个1x1卷积、以及一个identity直连每条路径后面都跟BatchNorm。三条路径的输出相加等价于让梯度有了多条回传通道同时相当于多个子网络在做隐式集成训练收益很直观。到了推理阶段这三条路径在数学上可以合并成一个3x3卷积。一旦合并完网络变成严格串行的3x3卷积堆叠没有跳跃连接没有并行分支。这种结构有两个工业级收益一是3x3卷积在绝大多数推理库TensorRT、OpenVINO、NCNN里都有最优算子的kernel匹配效率远高于1x1、5x5等异形核二是内存访问量大幅下降分支合并后中间特征不需要多路缓存Cache命中率会明显变好。简单说训练时我们在用算力换表达力推理时我们用代数换速度两头的优势都拿到了。这里要澄清一个最常见的误区重参数化不是剪枝不是量化的替代品而是把“结构里的冗余”在数值层面直接消掉。它和剪枝的差别在于剪枝会改变参数个数和网络拓扑重参数化不改变卷积核数量只是把原来由多个算子共同完成的计算合并成一个3x3卷积。正因为这个过程是纯线性的工程转换前后输出的差异理论上为零数值层面是浮点精度级别的差异这也让它的验证逻辑变得非常干净——转换前和转换后模型输出做余弦相似度比对几乎必过。2.2 单个RepVGG Block的实现细节实现一个RepVgg Block并不复杂核心是构造好三种分支并保证它们的输出可以按元素直接相加。下面是训练态Block的PyTorch实现我用最简单的方式把分支都列出来import torch import torch.nn as nn class RepVggBlock(nn.Module): def __init__(self, in_channels, out_channels, stride1, deployFalse): super().__init__() self.deploy deploy if deploy: # 推理态只有融合后的3x3卷积无BN无分支 self.rbr_reparam nn.Conv2d( in_channels, out_channels, kernel_size3, stridestride, padding1, biasTrue ) else: self.in_channels in_channels self.out_channels out_channels # 主分支3x3卷积 BN self.rbr_dense nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, stride, 1, biasFalse), nn.BatchNorm2d(out_channels) ) # 1x1卷积分支等价于小感受野特征提取 self.rbr_1x1 nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, 1, 0, biasFalse), nn.BatchNorm2d(out_channels) ) # identity分支需要保证inout且stride1 if in_channels out_channels and stride 1: self.rbr_identity nn.BatchNorm2d(in_channels) else: self.rbr_identity None def forward(self, x): if self.deploy: return self.rbr_reparam(x) if self.rbr_identity is None: return self.rbr_dense(x) self.rbr_1x1(x) return self.rbr_dense(x) self.rbr_1x1(x) self.rbr_identity(x)这段代码里有两个关键设计。第一identity分支并不是真正的“什么都不做”而是一个只对x做BatchNorm的直连这样在训练时它确实参与了归一化统计和数值加和但不会引入额外的卷积参数。第二1x1卷积和3x3卷积分支的输出形状必须一致否则无法相加所以当stride2时identity分支直接置空因为下采样后特征图尺寸已经对不上了。搭建完整网络时RepVgg按阶段划分不同版本的A0、A1、B1等主要区别在每阶段的Block数量和通道宽度。我建议第一次复现时先用RepVgg-A0跑通全流程它的参数量在几十MB级别单卡也能几小时内完成训练适合做逻辑验证。每个Stage之间用stride2的下采样接一个3x3卷积进行过渡分类头则是GlobalAvgPool加全连接层。2.3 重参数化的核心代数BN融合与三路合并这是整个RepVgg最值得自己推一遍的部分。训练结束后每个卷积层后面都跟着BN而BN在推理时完全可以用一组线性系数折算回卷积层里从而把“卷积BN”合并成一个带bias的卷积。这在数学上很直白BN对每个通道的归一化可以写作y gamma * (x - mean) / sqrt(var eps) beta而卷积的输出本身就是x所以把gamma、mean、var、beta折算进卷积核的weight和bias就行。下面的代码展示了如何把单个“卷积BN”融合成带bias的卷积这是所有重参数化工作的第一步def fuse_conv_bn(conv, bn): # 取出BN的统计量 gamma bn.weight.data beta bn.bias.data mean bn.running_mean.data var bn.running_var.data eps bn.eps # 计算缩放系数 t gamma / sqrt(var eps) t gamma / torch.sqrt(var eps) # 融合weight: 每个输出通道的卷积核按系数缩放 fused_weight conv.weight.data * t.view(-1, 1, 1, 1) # 融合bias: 如果原来卷积没有bias需要补上 if conv.bias is not None: fused_bias conv.bias.data * t beta - mean * t else: fused_bias beta - mean * t fused_conv nn.Conv2d( conv.in_channels, conv.out_channels, kernel_sizeconv.kernel_size, strideconv.stride, paddingconv.padding, biasTrue ) fused_conv.weight.data fused_weight fused_conv.bias.data fused_bias return fused_conv完成三路融合时还需要做两件事一是把1x1卷积pad成3x3尺寸在核矩阵中间放1x1的数值其余补0这样三个分支的核大小一致可以直接相加二是把identity分支构造为一个“单位矩阵型”的3x3卷积——中心位置为1、周围为0再跟着对上一个BN做融合。三个分支分别融合完后把weight和bias直接相加就得到最终单路3x3卷积。这一步最考验细心的地方是pad的坐标pad位置错了输出必然错位转换后的模型精度直接变成随机猜测。3. 在自定义数据集上跑通RepVgg图像分类3.1 数据集组织与训练配置图像分类任务的数据集组织其实就一个约定按类别建目录每个目录下放对应类别的图片。我最近在一个森林图像分类项目上就是这么做的按树种和病害类别分了7个文件夹训练集和验证集按8:2切分。用图像分类算法训练前先把目录结构切明确比写代码更影响最终效果。一个标准的PyTorch训练流程里数据加载部分需要固定的三件套读图、resize、归一化。RepVgg的这个阶段有个特别容易踩的坑——如果你的数据集图片分辨率不是224x224训练时强行resize会造成严重的信息丢失。我的习惯是先把数据集里图片的宽高分布跑一遍统计再看中位数决定resize目标。比如森林图像航拍图往往是几百像素到上千像素不等直接压到224会丢掉纹理细节我会先用短边resize到256再中心裁剪到224而不是简单resize。下面是一个典型的训练主循环骨架包含了模型创建和训练调度import torch import torch.optim as optim from torch.utils.data import DataLoader, Dataset from torchvision import transforms, datasets # 数据增强与归一化 transform_train transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.6, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) transform_val transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 构建model: 用RepVgg-A0num_classes换成自己的类别数 model build_repvgg_a0(num_classes7, deployFalse) optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100, eta_min1e-6) # 训练循环注意要设置model.train() for epoch in range(100): model.train() for batch in train_loader: x, y batch optimizer.zero_grad() out model(x) loss nn.CrossEntropyLoss()(out, y) loss.backward() optimizer.step() scheduler.step() # 验证逻辑省略注意要设置model.eval()这里有几个参数我不建议照抄就用。SGD的初始学习率0.01对应的是batch size 64到128左右如果你因为显存限制只能用batch size 16或者32学习率必须同步下调到0.001甚至更低否则前几个epoch的loss会剧烈震荡。CosineAnnealing的T_max要跟总epoch数对齐T_max设成50但实际训练100轮后半段的lr会走回头路收敛曲线会很难看。3.2 加载预训练权重做迁移学习从零训练RepVgg-A0在ImageNet尺度数据上大约需要两周时间在普通数据集上同样成本不低。实操里绝大多数情况应该用预训练权重做迁移学习。做法是加载ImageNet的预训练权重替换最后的全连接层节点数然后把整个网络的学习率整体调低。这里有个容易被忽略的细节重参数化模型在加载预训练权重时权重名里的rbr_dense、rbr_1x1、rbr_identity这些前缀必须严格对齐很多从官方仓库下载的权重用的key命名和自建模型的命名不同直接load_state_dict会报缺失键。遇到键名不匹配时不要用strictFalse硬跳因为那会让卷积层随机初始化训练精度会永久损失。正确的做法是把预训练权重字典打印出来把key名做映射后再加载。另外迁移学习时BN层的统计量也值得单独对待——如果新数据集的图像分布和ImageNet差异大比如红外图、医学影像、航拍森林前几个epoch要冻结BN的running_mean和running_var更新让模型先适应新分布再放开不然梯度会很快把统计量带入不稳定区间。3.3 训练时的显存控制与batch size选择RepVgg训练态因为同时存在三条分支显存占用比推理态高不少。A0模型在224分辨率下batch size 64需要约12GB显存如果是A1或B1模型会更吃紧。显存不够的常规解法是梯度累积它的原理很简单连续跑几个小batch每次累加梯度累积到目标batch大小后再更新一次参数。下面的代码展示了梯度累积的用法accum_steps 4 # 实际batch_size 显示batch_size x 4 for idx, (x, y) in enumerate(train_loader): out model(x) loss nn.CrossEntropyLoss()(out, y) loss loss / accum_steps # 把损失按累积步数归一化 loss.backward() if (idx 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()用梯度累积时BS和LR的匹配也要同步调整。假设显示batch size是32累积4步等效于实际的128那LR应该按128来配。另外别忘了同步调整BN的momentumBN的统计量是逐batch更新的小batch时统计噪声大建议把momentum从默认0.1调低到0.01让统计量变化更平滑。训练过程中最好隔一段时间就做一次“中途转换”把当前训练态的模型转成推理态在验证集上测一下准确率。如果训练态loss已经很低但推理态表现不对劲绝大多数情况是某个分支的BN统计量还没收敛可以再训练几个epoch。这个习惯能帮你提前发现问题比最后一次性转换失败再回头查省心得多。3.4 森林图像分类场景的实战配置以森林图像分类这个场景为例它的特殊性在于类内差异大、背景干扰强。你需要同时识别树种、枯死状态和火灾风险等级类别之间可能长得非常像。这时候普通的数据增强方案不一定够用我一般会加两部分一是随机擦除RandomErasing模拟树枝遮挡和烟雾模糊二是CutMix或MixUp让模型对混合背景更鲁棒。另一个关键点是类别不平衡问题。森林巡检数据里健康树木的数量往往远多于病害树直接训练会让模型把所有样本都预测成多数类。解决的办法是在CrossEntropyLoss里传入一个按类别频次倒数归一化的weight向量。计算方式不复杂统计每个类别的样本数取倒数后归一化。这样少样本类别的loss权重更大模型会花更多注意力在难分类别上。数据的归一化参数也建议根据森林图像重新统计而不是沿用ImageNet的均值方差。因为航拍森林图像的通道分布和日常照片差别很大绿色通道的均值会明显偏高用ImageNet统计量强行归一化等于把特征分布做了错误的平移。虽然模型最终也能训练出结果但前期收敛会慢很多尤其在用预训练权重做迁移时分布错配的影响会更明显。4. 推理部署模型重参数化转换与四个血泪坑4.1 手写转换脚本从训练态到推理态一键融合重参数化转换是整个RepVgg工程里最不能出错的一环。转换的基本流程是遍历每个Block对三条分支分别做conv-BN融合然后把1x1和identity分支重排成3x3核三核相加得到最终参数。下面是一个可以直接套用的转换函数按Block粒度处理def reparametrize_model(model): for module in model.modules(): if isinstance(module, RepVggBlock): # 三路卷积核提取 # 第一步把dense分支(3x3)与BN融合 fused_dense_w, fused_dense_b fuse_conv_bn( module.rbr_dense[0], module.rbr_dense[1] ) # 第二步把1x1分支pad成3x3 fused_1x1_w, fused_1x1_b fuse_conv_bn( module.rbr_1x1[0], module.rbr_1x1[1] ) padded_1x1_w torch.zeros(fused_1x1_w.shape) padded_1x1_w[:, :, 1:2, 1:2] fused_1x1_w # 第三步identity分支构造3x3单位卷积 identity_conv nn.Conv2d( module.in_channels, module.out_channels, kernel_size3, stride1, padding1, biasTrue ) identity_conv.weight.data.zero_() identity_conv.weight.data[:, :, 1:2, 1:2] torch.eye(module.in_channels) fused_id_w, fused_id_b fuse_conv_bn(identity_conv, module.rbr_identity) # 三路相加 final_w fused_dense_w padded_1x1_w fused_id_w final_b fused_dense_b padded_1x1_b fused_id_b # 用单路卷积替换原Block module.rbr_reparam nn.Conv2d( module.in_channels, module.out_channels, kernel_size3, stridemodule.rbr_dense[0].stride, paddingmodule.rbr_dense[0].padding, biasTrue ) module.rbr_reparam.weight.data final_w module.rbr_reparam.bias.data final_b # 删除训练用分支 module.rbr_dense None module.rbr_1x1 None module.rbr_identity None module.deploy True return model这段代码最值得注意的是pad 1x1卷积时的索引[1:2, 1:2]正好是3x3核的中心位置。很多人会把1:2写成1那样会少一个维度导致赋值报错。另外identity分支融合时构造的“单位卷积”权重中心为1、四周为0含义是输出等于输入这是代数恒等映射的卷积化表达。转换完成后模型里只剩下一串3x3卷积加激活函数后续可以接TorchScript导出、ONNX导出或NCNN转换。ONNX导出时有一个额外好处因为网络里没有分支和跳跃连接ONNX的自动优化几乎不需要做额外工作导出的图结构就非常干净。4.2 转换后验证相似度比对数类结果转换完成后绝对不能直接上部署第一步是验证转换的正确性。最可靠且简单的方式是取同一批输入分别输入转换前的训练态模型和转换后的推理态模型比较输出特征图或最终分类logits的差异。理论上两种模型输出应当完全一致实际因为浮点累加顺序不同会有1e-5到1e-4量级的微小差异这完全正常。如果差异超过1e-2转换脚本一定有bug。def verify_reparam(model_train, model_deploy, sample_input): model_train.eval() model_deploy.eval() with torch.no_grad(): out_train model_train(sample_input) out_deploy model_deploy(sample_input) # 计算余弦相似度数值越接近1越好 cos torch.nn.functional.cosine_similarity( out_train.view(1, -1), out_deploy.view(1, -1) ) print(f余弦相似度: {cos.item():.6f}) # 计算最大绝对误差 max_diff (out_train - out_deploy).abs().max().item() print(f最大误差: {max_diff:.2e})验证时必须在model.eval()状态下进行。如果在train模式下跑转换前的模型BN还在用当前batch的统计量更新两组输出肯定对不上这不是转换错了而是模式错了。这个错误我见过太多次十次检查里有五次是这个原因。另外验证用的输入最好多找几类有代表性的图正常光线的、暗光下的、带噪声的确保不是只在某一张图上数值恰好吻合。4.3 坑一BN在train和eval模式下的统计量差异导致误判这是新手最容易误判的点。训练态模型在train模式下BN会实时更新running_mean和running_var但转换过程用到的是BN最终保存好的running_mean和running_var。如果转换时训练态模型的BN统计量还没稳定或者你在转换前跑了几步train模式前向统计量就会被污染。解决方法是转换前专门调用一次model.eval()并在验证集上跑一个epoch让BN统计量“稳定落地”再做转换。更隐蔽的情况是使用分布式训练或DDP时每个卡上的BN统计量在最后的global平均之前可能与单卡推理不一致。多卡训练完保存的权重在单卡加载后跑evalBN的统计量是全局平均过的这是好事但如果你在多卡模式下直接做模型转换并校验每个rank的输出会因为不同rank数据分布不同而产生误差。这类问题在单卡部署时通常不存在一旦上多机TensorRT推理就可能冒出来。稳妥做法是转换和校验全部在单卡CPU或单GPU环境下完成。4.4 坑二自定义部署代码里误用1x1卷积算子转换完的推理态网络如果还有人试图从结构上去找1x1卷积分支那是白费力气因为已经不存在了。真正的坑出现在目标部署框架里有些推理框架的卷积算子实现不够全面如果代码里残留了对非3x3卷积的调用比如你写了一个通用的Block类支持任意kernel size框架可能选不到最优kernel推理速度会打折扣但精度没问题。更危险的情况是某些框架尤其是移动端框架会对3x3卷积做Winograd变换以加速而Winograd在特定输入尺寸下需要padding对齐比如NCNN有pad到4的倍数约束。RepVgg网络如果中间有非3x3的全局池化或其他异形操作整个图会被多次切分Winograd优势会缩小。所以部署前一定要看清推理框架的算子约束如果条件允许把输入分辨率固定成框架友好的尺寸能省掉很多运行时判断的浪费。4.5 坑三数据预处理不一致导致推理精度变成随机猜测转换后模型结构完全正确但推理精度崩成50%甚至个位数这是最常见的部署翻车。原因几乎都在预处理训练时用的是Resize(256) CenterCrop(224)部署端却直接resize到224x224或者训练时的归一化用的ImageNet均值方差部署端写错了std的顺序。这些不一致会直接把输入特征分布推到模型从未见过的区域。排查这类问题有个很笨但有效的办法在部署环境里直接把一张验证集图片按部署流程做预处理再打印出预处理的numpy结果和训练脚本里同一张图的预处理结果做逐元素相减看差异是否为0。这个方法可以在十分钟内定位绝大多数预处理Bug。真正落地的项目里我一般会把预处理写成单独的函数训练端和部署端共用同一个函数从根源上避免两套逻辑分叉。4.6 坑四模型转换后部署端运算精度从FP32降到FP16导致的掉点RepVgg转换完的模型在FP16下推理时偶尔会出现精度掉1到2个点的现象。原因是重参数化后所有参数融合进单个卷积核数值范围可能比原来大尤其是BN里的gamma如果比较大会让融合后的权重数值范围变大FP16的动态范围扛不住。解决方案有三个方向第一转换后检查权重数值分布如果最大值超过FP16的合理范围考虑在BN融合后对权重做一次max-norm缩放但这会破坏原模型输出等价性需要重新微调第二做混合精度推理时只对激活值用FP16权重保留FP32第三使用INT8量化前先做一次“校准”让模型重新适应量化分布不要直接拿FP32权重硬转。这些都是工程上的细活我的建议是先用FP32把整个精度基准确立下来再动量化。5. 把精度和速度一起调的三个验证手法和一个习惯模型转换成功后工作并没有结束。真正决定这个方案值不值得投入的是三个维度的验证一是精度是否无损二是推理速度是否真比经典模型快三是增益在不同硬件上是否稳定存在。我习惯用一套固定的对比流程来做这件事。精度验证除了前面说的余弦相似度外还要在完整验证集上跑一次Top-1准确率对比。做法是把训练态模型和推理态模型分别在同一条验证pipeline上评估对比两个准确率应该完全一致。这里有个容易被忽略的统计问题如果验证集比较小两个模型的输出差异虽然只有1e-4量级但接近决策边界的样本可能在阈值两侧抖动导致准确率差0.1-0.2个百分点。这种情况下我直接把两份logits做argmax后再统计不一致的样本数量如果只有极少数且置信度都在0.5附近说明数值差异在正常范围内不用紧张。速度验证要用同一份代码、同一块硬件、同一个batch size分别测RepVgg转换后模型和一个对照模型比如ResNet50的FP32以及FP16推理时延。不要只看端到端时间还要用profiler看卷积层耗时占比、显存占用和Cache命中情况。RepVgg强项是卷积层密集且计算率恒定所以当你的硬件3x3算子优于其他形状算子时收益会特别明显如果目标硬件对3x3有特殊限制建议先用benchmark程序跑一下方差再决定。最后一个习惯也是我每次做图像分类实战都会坚持的初期用小模型做完整流程验证再放大模型。也就是先用RepVgg-A0把数据、训练、转换、部署全链路的正确性确认下来然后才切换到A1或B1去刷精度。这样做的理由很简单——重参数化的坑和模型大小无关小模型踩完坑大模型再踩的概率就会低很多。我曾经为了省事直接训练B1模型结果转换后精度异常排查了两天才发现是数据预处理里std顺序写反了这个错误在A0上只要两小时就能定位出来。之后我淡定多了凡是接入新数据集、新部署框架一定用小模型跑一个端到端全流程再谈放大训练。RepVgg这类“训练重参数、推理单路径”的思路在图像分类算法这条路上还会延续很久因为它不挑硬件、不挑框架、不挑开发语言只需要你在部署前多花一点时间写转换代码。方向值得投入成本主要在前期流程验证一旦把转换和校验固化到工程流程里后面每次换数据换设备都会觉得很顺手。希望这些踩过的坑和验证思路能帮到你像当时帮我一样少走几段弯路。本文还有配套的精品资源点击获取
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →