ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

YOLOv8s通道剪枝实战:从稀疏化训练到边缘部署加速

YOLOv8s通道剪枝实战:从稀疏化训练到边缘部署加速 先说结论剪枝不是玄学它本质上是先把模型里不重要的通道找出来删掉再花一小段时间把精度养回来的工程流程。这篇文章我会把YOLOv8s从稀疏化训练、通道筛选、模型重建到微调的完整源码思路拆开讲相关代码可以直接改改路径拿来用。我手头这个项目原本是用YOLOv8s做产线缺陷检测训练机是两张RTX 3090精度没问题但真正落地时发现推理机只有一张老旧的Tesla P48GB显存勉强装得下跑起来却只有十几帧。换轻量模型精度又不够最后决定走剪枝路线。这次实操从准备数据到剪枝完成、精度回升前后花了一周多中间踩了不少坑这篇就当是给同样卡在部署环节的同学一份带源码的参考手册。1. 剪枝前先摸清YOLOv8s的家底参数都藏在哪任何剪枝工作开始之前都得先搞清楚模型的钱花在什么地方。YOLOv8s整体参数量大约11.2MFLOPs约28.6G对边缘设备来说算不上特别重但也没到随便跑的地步。它的结构由三块组成CSPDarknet风格的Backbone负责提特征PAN-FPN结构的Neck负责多尺度融合Decoupled Head负责分类和回归。用一段很简单的代码就能统计各部分参数量model DetectModel(cfgyolov8s.yaml, ch3, nc80) total 0 for name, module in model.named_children(): params sum(p.numel() for p in module.parameters()) print(f{name}: {params / 1e6:.2f}M) total params print(fTotal: {total / 1e6:.2f}M)实测下来Backbone占掉总参数量的55%左右Neck约占30%Head约15%。如果你的业务和我的类似是缺陷检测这类对边缘设备部署敏感的任务那剪枝的重心就应该放在Backbone和Neck上Head的通道数受输出维度约束自由度反而不大。剪枝前还要理解YOLOv8s里每个模块的结构。C2f模块是从CSPNet进化来的它内部先经过一个Conv投影然后拆成两条路径一条直接连后续模块另一条经过n个Bottleneck堆叠最后再Concat在一起。这种结构在剪枝时有个麻烦如果只把某个中间层剪了Concat之后的通道数就对不上会导致shape错误。所以剪枝的时候要么按照组来做要么干脆按照C2f整体输出通道数统一处理。后面我会讲到我的处理方式。另一个关键点是BN层的分布。YOLOv8默认在每个卷积后面都接BN层而BN层的γ缩放系数直接决定了这个通道对后续特征的影响强度。这就是通道剪枝的理论基础γ值小到接近0的通道意味着它输出的特征图经过缩放后对最终结果几乎不起作用删掉它们是安全的。在动手前我习惯先用TorchScript把模型导出来跑一遍随机输入确认原始mAP基线。这一步不是为了炫技而是为了给剪枝后的对比留一个手术前体检报告。我这里把基线定义为mAP0.5:0.95约44.8mAP0.5约62。后面对比就围绕这两个数来谈。2. 结构化剪枝与稀疏化训练为什么我选了通道剪枝这条路剪枝大体分两种流派非结构化剪枝和结构化剪枝。非结构化剪枝是把单个权重置零稀疏度可以拉得很高但实际推理时权重矩阵仍然是稠密存储的如果没有配套的稀疏算子库基本跑不出加速效果——纯属纸面压缩。PyTorch自带的torch.nn.utils.prune就是这类它可以很方便地把权重剪成稀疏矩阵但对CNN在GPU上的实际推理速度几乎没有正向帮助。结构化剪枝则是以通道为单位整条删掉模型结构会真正变窄宽度缩小后FLOPs和访存量同时降低TensorRT、OpenVINO这类推理引擎都能直接吃到收益。通道剪枝有几个具体实现思路基于BN层γ稀疏化把γ推向0再按阈值删除通道基于特征图秩分析计算每个通道输出特征图的秩删掉低秩通道基于泰勒展开估计每个通道对损失的贡献贡献小的先删基于SLRSparse-Low-Rank分解将权重分解后裁掉奇异值小的部分。在YOLOv8s上最稳、源码资料最多的还是BN层γ稀疏化。它的核心思路其实很朴素给BN层的γ参数加上L1正则化让γ在训练过程中自然地向0收敛。训练完了之后统计所有γ值的分布设定一个全局阈值低于阈值的通道直接删除再重建一个小号模型。为什么L1正则能把γ推向0而不是像L2那样只把数值变小因为L1的梯度是常数符号每次更新都会给γ一个固定的向零拖拽力而L2的梯度跟γ本身大小成正比γ已经很小时拖拽力也变小很难真正归零。这个过程可以类比成砍树和拧螺丝的区别L1是对每棵树持续锯根部L2是在拧一颗已经快松的螺丝越拧越费劲还不一定拧得动。稀疏化训练的损失函数长这样loss main_loss s * sum(torch.abs(m.weight).sum() for m in model.modules() if isinstance(m, nn.BatchNorm2d))其中比例因子s是要调的。s太小γ稀疏得慢剪枝时找不到足够多的低值通道s太大主任务的损失被带偏mAP掉得厉害。实践里我一般从1e-4起步观察两三个epoch的γ直方图再微调。这一步是整个剪枝流程的天花板稀疏化质量直接决定后续剪枝能剪多狠。在正式稀疏化训练之前我还额外做了一件事把训练分成两个阶段。第一阶段先用原始损失训练若干epoch让模型回到合理的精度区间第二阶段才把稀疏化损失加进去。为什么这么做如果从一开始就加正则模型在训练的早期阶段就会过度牺牲主干能力去向0收敛后面很难通过微调拉回来。这有点像跑步前先热身再上强度道理很简单但很多人图省事直接一气呵成结果精度一塌糊涂。3. 稀疏化训练的实现姿势损失注入、训练策略与γ分布观察讲具体实现。YOLOv8s由Ultralytics框架托管训练直接改源码里的损失函数是最省事的做法。如果你用的是官方仓库稀疏化损失需要挂到每一个BatchNorm2d上并且要在每个step里叠加到总loss上。def add_sparsity_loss(model, loss, s1e-4): sparse_reg 0.0 for m in model.modules(): if isinstance(m, nn.BatchNorm2d): sparse_reg torch.abs(m.weight).sum() return loss s * sparse_reg把它接到训练循环里在每个step计算完原始loss之后调用。要注意的是官方训练脚本用的是AMP混合精度而torch.abs().sum()这类操作在FP16下没问题但累加到主loss时要注意精度被截断的问题。我自己是直接把稀疏项改成sparse_reg.float()再相加避免梯度溢出。稀疏化训练的scheduler也要做调整。官方YOLOv8s训练默认从lr0.01开始我用余弦退火跑300个epoch。在实际操作中建议先跑前50个epoch不加稀疏损失之后再加。这个预热阶段的目的我前面提了是为了让模型先稳住基本特征提取能力。后面的稀疏化阶段s值可以保持恒定也可以逐步增大我个人经验是恒定的效果比递增好调参。训练过程中要实时看γ分布。每跑完一个epoch把模型里所有BN层的weight收集起来画一个直方图。直观判断标准是正常训练的γ分布大概是一个集中在0.5到1.5之间的钟形而稀疏化训练后期应该出现一个明显的高峰聚集在0附近同时还有一小撮γ保持较大值。如果看到这个两根触角的形状就说明哪些通道值得保留已经写在了γ值里。import torch import matplotlib.pyplot as plt bn_weights [] for m in model.modules(): if isinstance(m, torch.nn.BatchNorm2d): bn_weights.extend(m.weight.detach().cpu().numpy().tolist()) plt.hist(bn_weights, bins200, range(-0.5, 2.0)) plt.savefig(fgamma_hist_epoch_{epoch}.png)这里有一个容易栽进去的坑稀疏化训练结束后很多γ值确实收敛到了0附近但并不会精确等于0。剪枝时千万不要直接判断γ 0而是要用一个阈值去截断。阈值怎么定常见做法是把所有γ绝对值排序取第P个百分位作为阈值P就是你想剪掉的通道比例。比如想让模型缩水50%那就取γ绝对值的第50百分位数作为阈值低于它的通道全删。我在实操中发现YOLOv8s的BN层数量比普通ResNet要多不少尤其是C2f内部的Bottleneck每个都有BN层导致γ的分布混合了不同层级的统计特性。直接拿全局百分位一刀切有时会误伤某些原本很重要的层。我的补救方案是对每个模块分别计算阈值同时加一个全局约束来限制整体删除比例。这比一刀切稳妥尤其是对C2f这种内部有短路连接的结构。4. 通道筛选与模型重建核心源码拆解与维度对齐模型重建是整个剪枝过程里代码最绕、也最容易出错的一环。PyTorch没有提供现成的删掉某个通道的APItorch.nn.utils.prune只能置零权重不能改变权重矩阵的shape。所以必须做的事是根据γ掩码重新构建一个小宽度的模型再把原模型的权重拷贝过去。我把这个过程分四步走。4.1 生成剪枝掩码def generate_mask(model, prune_ratio0.5): bn_weights [] for name, m in model.named_modules(): if isinstance(m, torch.nn.BatchNorm2d): bn_weights.append(m.weight.detach().abs().cpu().numpy()) all_bn np.concatenate([w.flatten() for w in bn_weights]) threshold np.percentile(all_bn, prune_ratio * 100) masks {} for name, m in model.named_modules(): if isinstance(m, torch.nn.BatchNorm2d): weights m.weight.detach().abs() mask weights threshold # 防止某个BN层所有通道全被剪掉 if mask.sum() 0: mask[weights.argmax()] True masks[name] mask.cpu().numpy() return masks这个masks字典记录了哪些采样点索引的通道应该保留。后面重建模型时每一层的输入输出维度都要根据上一层的mask计算出来。注意mask生成时有个关键约束——首层Backbone第一层卷积和Detect输出层不能剪。第一层卷积的输入是3通道RGB你不能把输入通道剪了Detect头输出维度由类别数决定更不能动。这个在生成mask后要手动排除。4.2 逐层重建模型结构这里有两种路线。路线一是改配置文件yolov8s.yaml把每个层的宽度重新算好再用Ultralytics去构建新模型然后手动拷贝参数。路线二是写一个通用的模型replacer动态遍历原模型遇到Conv和BN就替换成新尺寸的层。我推荐路线二因为YOLOv8的模型定义里有很多Concat、Shortcut这类形状依赖操作改YAML很容易漏掉细节动态遍历至少能把所有层都在内存中跑一遍。def slim_model(model, masks): # 先记录每层原模块类型与必要属性 new_modules {} prev_channels 3 for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): mask masks.get(prev_bn_name) if mask is not None: new_in_channels int(mask.sum()) else: new_in_channels prev_channels new_out_channels int(module.out_channels) new_conv torch.nn.Conv2d( new_in_channels, new_out_channels, module.kernel_size, module.stride, module.padding, groupsmodule.groups, biasmodule.bias is not None ) new_modules[name] new_conv prev_channels new_out_channels elif isinstance(module, torch.nn.BatchNorm2d): mask masks.get(name) if mask is not None: new_bn torch.nn.BatchNorm2d(int(mask.sum())) else: new_bn torch.nn.BatchNorm2d(module.num_features) new_modules[name] new_bn # 其他层激活、shortcut等按原样复制 return new_modules这只是一个示意轮廓实际处理要复杂得多。最大的坑在于Conv和BN是成对出现的但YOLOv8的命名里Conv模块内部包了一个Sequential里面既有Conv2d又有BatchNorm2d你直接遍历named_modules()时model.conv和model.bn是分开暴露的这说明你需要在module层面做整体替换而不是一层层地单独替换。我在源码解析时会用一个RepConv风格的替换逻辑遇到一个包含ConvBNSiLU的复合模块整体计算出新模块的输入输出通道然后一次性替换。4.3 权重参数搬运模型骨架重建好之后剩下是权重搬运。核心逻辑是根据mask保留对应索引的权重。卷积核的形状是[out_ch, in_ch/groups, kh, kw]那么def copy_weights(old_conv, new_conv, in_mask, out_mask): # old_conv.weight shape: [old_out, old_in, kh, kw] if in_mask is not None: w old_conv.weight[:, in_mask, :, :] else: w old_conv.weight if out_mask is not None: w w[out_mask, :, :, :] new_conv.weight.data.copy_(w) if old_conv.bias is not None: if out_mask is not None: new_conv.bias.data.copy_(old_conv.bias[out_mask]) else: new_conv.bias.data.copy_(old_conv.bias)对应BN层def copy_bn(old_bn, new_bn, mask): new_bn.weight.data.copy_(old_bn.weight[mask]) new_bn.bias.data.copy_(old_bn.bias[mask]) new_bn.running_mean.data.copy_(old_bn.running_mean[mask]) new_bn.running_var.data.copy_(old_bn.running_var[mask])这里特别要注意的是如果模型里有Concat操作Concat的两个分支输入通道必须保证一致。YOLOv8s的Neck就是这种典型结构C2f的输出和Upsample出来的特征要拼在一起这两条路径的通道在剪枝时必须按同一个mask裁剪不然拼接时就会报错。我的解决方式是先构建一张通道传播依赖图把因为Concat而耦合的层标记成同一剪枝组整个组共享一个mask。这个方法实现时有几个细节C2f内部的shortcut分支也要跟着主分支剪否则残差对齐不了。4.4 验证新模型输出shape权重搬运完后用随机张量跑一遍前向确认没有shape报错。这一步不能偷懒。我实际遇到过看似每一步都对齐了但某个SPPF层输出的通道没更新导致Detect头输入和输出维度对不上前向时直接炸掉。把前向跑通后再导入一批验证集图像算一次剪枝后的mAP看看在没有微调的情况下掉了多少。如果剪枝率设置合理比如50%剪完后直接推理的mAP通常会比原模型掉10到15个百分点。这个回落幅度是正常的千万别觉得模型废了。只要没掉到个位数微调阶段都有机会拉回来。如果掉得离谱优先检查是不是某个关键层被误剪了。5. 剪枝后微调的3个关键设置学习率、训练轮数与蒸馏辅助微调阶段直接决定剪枝最终效果这里的经验含量比前几步更高。第一是学习率。剪枝后的模型是一个受伤的模型它原本的损失表面被强行改变了形状。如果沿用原始训练的高学习率很容易在一次更新里把权重踢到很差的局部极小值。我建议把学习率设置为原训练最后阶段学习率的1/10左右。比如原来用0.01的余弦退火那微调和退火初始值就放在0.001并且用warmup把前几个epoch推上去再退火。第二是训练轮数。剪枝模型通常不需要从头训300个epoch。因为这模型不是从随机初始化开始的它有绝大部分原模型权重作为底子。我实测过在同样数据集上微调80到120个epoch就可以把mAP拉回到剪枝前的95%以上再往后训练收益越来越小。轮数太多还会过拟合到训练集在验证集上的收益反而下降。第三是蒸馏辅助。如果你的原模型还在可以顺手把原模型当teacher把剪枝后的模型当student在微调阶段加一个蒸馏loss。蒸馏的方式有两种硬蒸馏直接用原模型的预测logits当soft label去训练student特征蒸馏让student的中间特征图去匹配teacher的特征图。考虑到YOLOv8的head输出是解耦的我建议用硬蒸馏就好了它对项目侵入最小实现成本也最低。在微调阶段总loss写成loss cls_loss box_loss dfl_loss lambda * distill_loss其中distill_loss用KL散度或者简单的MSE让student输出的类别概率分布逼近teacher。这里还要注意一个细节稀疏化训练用的s参数在微调阶段必须关掉否则模型好不容易拉回精度又被正则项往0拽。很多人忘了这一步结果微调完整整十几个epoch精度纹丝不动。6. 从YOLOv8s剪到YOLOv8m量级的实测记录数据与经验我用自己的数据集4500张缺陷样本分8类完整跑了一遍。数据集不大但分布很典型小目标多、背景杂YOLOv8s原本在验证集上能做到mAP0.5约84.6、mAP0.5:0.95约58.3。稀疏化训练阶段跑120个epoch前30个epoch预热后90个epoch加入L1正则s2e-4。训练结束后γ分布的高峰明显聚集在0.05附近。用第50百分位做阈值剪掉约52%的通道模型宽度统计对比模块区域剪枝前通道数剪枝后通道数压缩比Backbone前段32/64/12816/32/64约50%Backbone后段256/512128/256约50%Neck融合层256/128/64128/64/32约50%Detect头64/12864/128不剪剪完直接推理mAP0.5从84.6掉到71.2mAP0.5:0.95掉到44.7。掉得确实很狠但也在预期内。然后进入微调阶段学习率从0.001开始余弦退火跑100个epoch。最终结果mAP0.5恢复到82.3mAP0.5:0.95恢复到57.1只比原始模型低了不到1.5个点。参数量的变化更可观。原模型11.2M剪完只剩4.8M降幅约57%。FLOPs从28.6G降到约12G左右。在没有用TensorRT、纯PyTorch推理的情况下单张Tesla P4上的推理耗时从62ms降低到33ms帧率从16提升到30。如果换成TensorRT INT8量化还能再快一截。用一句话形容就是花7天时间换来了一次模型减肥成功。这组数据放在更大模型上同样有参考性。如果你用的是YOLOv8m/l/x剪枝的收益理论上更大但要注意模型越大训练数据量必须越充足否则微调阶段的恢复力不够。7. 实操中遇到的坑和我的应对从shape崩溃到mAP塌方剪枝这个活坑藏在细节里。我把自己踩过的大坑和对应的排查思路写出来希望后来的人少走点弯路。第一个坑是C2f模块的shape错乱。YOLOv8的C2f内部做了split和concat剪枝后数据流在模块内部对不上前向时直接报错。排查链路我建议这样走先用最小的网络把配置里的depth和width改成很小的值跑一遍重建逻辑确认思路没有结构性问题再上真实配置。如果你连最小配置都跑不通说明你的重建代码没有正确处理模块内部的通道依赖这时候不要急着去调大模型停下来修改重建逻辑。第二个坑是剪枝率设得太高。一开始我图快设置了70%的剪枝率结果微调50个epochmAP只恢复到64%。后来我降回50%同样的微调流程恢复到了82%。这个对比非常直观剪枝率超过某个临界点之后模型的可恢复性会断崖式下降。我的经验是对YOLOv8s这类中小型模型50%到55%是一个比较稳的目标区间想剪到70%以上就需要配合蒸馏和更长的微调周期。第三个坑是epoch轮的稀疏化训练过度。有些前辈的教程会告诉你稀疏化训练可以跑到全部epoch都带正则我试过之后发现效果并不好。原因很简单模型在训练后期本来就在收敛这时再强推γ向0会严重破坏已经学好的特征空间。所以我的建议是稀疏化训练的后半程就够用最多不要超过总epoch的2/3。第四个坑是BN层的running_mean和running_var处理。剪枝时很多人只copy了weight和bias忘了copy统计量。这会导致模型在BatchNorm层输出的分布错乱推理mAP异常地差。这个坑非常隐蔽因为前向不会报错数字看起来只是差不会让你想到是BN统计量的问题。我在代码里写了一行注释copy BN running stats too, otherwise distribution shifts silently.第五个坑是Detect头的通道约束。YOLOv8的Detect头输出通道是nc 4的衍生维度分为分类分支和回归分支。如果剪枝的mask不小心作用到了Detect头内部的Conv层输出通道数会直接不符合类别数的要求推理时报错或者输出张量形状错误。所以Detect头要么整体跳过要么只对Backbone送入的特征做裁剪但保持内部输出通道不变。第六个坑是数据加载时的显存溢出。剪枝后的模型虽然显存占用降低了但微调阶段如果沿用原始的batch_size显存会多出来一截。这时候不要盲目调大batch_size因为BatchNorm的统计量在剪枝后需要重新适应分布batch越大反而会让BN统计更新过于平滑不利于快速恢复。我微调时把batch_size保持和原训练一致或者只上调一点效果好于大幅加大batch。8. 把剪枝流程工程化的建议脚本分层与参数排查剪枝不是一次性跑完就结束的事。项目如果还要迭代建议把流程拆成几个独立的脚本各自负责一个阶段避免每次改一个参数都从头跑一遍。我的目录结构大概是这样的prune_yolov8s/ ├── 01_sparse_train.py # 稀疏化训练入口 ├── 02_analyze_gamma.py # 生成γ直方图和阈值分析 ├── 03_prune_rebuild.py # 生成mask、重建模型、拷贝权重 ├── 04_finetune.py # 剪枝后微调 ├── 05_evaluate.py # 统一评估mAP和速度 └── configs/ └── prune_config.yaml # 存放剪枝率、s、阈值、epoch等每个脚本都从同一个config.yaml读参数这样改一个参数只需改配置不用动代码。对于团队协作这个习惯能避免很多你们怎么跑出来的数字我不对的扯皮。剪枝率这类关键参数我建议在config里写全评估记录。比如把每次运行的剪枝率、剪后mAP、微调后mAP、推理耗时都记录下来横向对比。如果同一个剪枝率重复跑出偏差很大的结果先查随机种子和数据加载顺序再查是否误改了稀疏化阶段epoch数。关于数值稳定性再补一句稀疏化训练阶段最好固定随机种子。我自己固定了torch.manual_seed(42)和np.random.seed(42)锁定数据增强的随机序列。剪枝对随机性比普通训练敏感得多不固定种子的话同样的config跑两次γ分布都会有可见差异剪出来的模型结构都可能不一样。另外我强烈建议把剪枝后的模型导出成ONNX和TensorRT引擎各验证一次。ONNX导出能发现动态shape问题TensorRT能直观看到推理加速。YOLOv8官方仓库已经支持导出但剪枝后的自定义模型结构需要自己处理导出时的opset兼容性。我踩过的一个小坑是nn.Upsample在导出时用了nearest模式TensorRT里要选对对齐方式否则特征图有像素级偏移mAP会莫名其妙掉一截。项目走到收尾之后我最大的体会是剪枝的成功率7成取决于前期稀疏化训练的质量2成在微调1成才是那个剪枝算法本身。很多人在网上找一键剪枝脚本拿过来跑不出理想效果根源往往不是剪枝源码有问题而是他们跳过了稀疏化这一整块前置工作。剪枝不是给训练好的模型做手术而是从训练阶段就为之后的瘦身打好麻药、做好标记。理解了这一点再看别人的剪枝源码思路就清楚了。
返回列表