ARTICLE DETAIL

资讯详情

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

YOLOv8s结构化剪枝实战:从BN稀疏化到模型通道重建

YOLOv8s结构化剪枝实战:从BN稀疏化到模型通道重建 做模型压缩这两年我经手的检测模型里YOLOv8s剪枝算是最常被问到的需求之一。原因很简单YOLOv8s本身是个不错的基线但部署到边缘设备或者追求高帧率时参数和FLOPs还是嫌多。剪枝尤其是结构化剪枝是比量化更温和、也更考验细节的一条路。今天我就把这套针对yolov8s的剪枝源码拆开讲清楚从原理到踩坑顺序捋一遍你需要跑通或者复现的时候按这个来就行。1. 项目概述与需求拆解1.1 为什么选择yolov8s作为剪枝对象YOLOv8s在YOLOv8系列里属于“小但够用”的定位。相比n版本精度更高相比m版本速度快不少很多项目拿它当初始模型训练完直接转ONNX再用TensorRT部署。但问题也随之而来如果你的摄像头路数多、芯片算力有限或者对延迟有硬性要求哪怕是s版本也得再瘦身。剪枝的目标不是把模型压到极致而是尽量在保住mAP的前提下把冗余通道去掉。我做过几次实验物体检测模型经常存在大量接近零贡献的通道尤其是骨干网络后几层的特征通道。这些通道在训练时参与了前向传播但重要性很低删掉之后网络依然能输出相近的结果。yolov8s因为结构规整、卷积层密集剪枝效果好也容易分析不像一些轻量化网络本身通道已经抠得很紧剪了就崩。1.2 剪枝能带来什么实际收益大家最关心的其实是三个数字模型体积、FLOPs、推理延迟。对一个输入尺寸640x640的yolov8s来说原始模型大小大概在22MB左右权重文件FLOPs在28G左右。我用这套源码做过一轮实验剪掉35%的通道后模型体积降到12MB上下FLOPs降到16G附近GPU推理延迟约降低了25%CPU上降低更明显。不过这里我要先泼一盆冷水延迟不一定和FLOPs严格同比降低。原因在于剪枝后的通道数如果不满足某些硬件对齐要求内存访问的连续性反而受影响。TensorRT优化之后效果会好很多但直接在PyTorch里测有时候只快10%到15%。所以你在看效果时一定要用目标推理框架测不要只盯着FLOPs数字。1.3 这套源码到底解决什么问题网上的剪枝代码零零散散很多是针对分类网络的拿到检测模型上直接跑不通。YOLOv8有C2f模块和Detect检测头结构比ResNet复杂shortcut连接也更多。如果直接套用普通通道剪枝脚本弹出的第一个错误就是维度不匹配。这套源码的核心思路是围绕YOLOv8本身定制的主要解决三件事支持对C2f模块、骨干和检测头做通道级剪枝自动处理shortcut残差分支的通道对齐问题提供剪枝后的模型结构重新生成策略避免手工改yaml改到崩溃简而言之你需要的不是一个通用剪枝库而是一套能适配实际检测模型、能直接落地跑通的流程。下面的内容就是把这套流程的每个环节说清楚。2. 剪枝方案选型结构化剪枝与非结构化剪枝2.1 两种主要剪枝路线对比剪枝算法的大致分类很多新手一开始会混淆。简单说非结构化剪枝是把权重矩阵里接近0的小权重单独置零得到的是稀疏矩阵。这种剪枝不会改变网络结构但需要专门的稀疏推理库或者硬件支持才能加速。所谓“细粒度剪枝”说的也是这种。非结构化剪枝的压缩率可以做得很高但在常规CPU和GPU上没有明显加速收益。结构化剪枝是把整个通道或者整个卷积核删掉。删除之后特征图的通道数变少后续卷积层的输入维度也变小网络结构真实变瘦。这种剪枝不需要特殊硬件任何深度学习框架都能正常加载推理。我说得直白一点你要部署选结构化剪枝你只是研究算法或者追求极限压缩率非结构化剪枝可以玩但工程化价值低很多。这套源码做的通道剪枝就是结构化剪枝的一种特别适合YOLOv8这种多卷积层的网络。2.2 为什么选择基于BN层的通道剪枝通道剪枝的关键问题是怎么判断哪个通道该剪。判断方式有很多比如基于权重大小、基于激活值统计、基于梯度信息但工程上最常用也最好实现的是借助批归一化层BN的缩放因子gamma。YOLOv8的每个卷积模块后面基本都跟了BN层。BN做的事情是把输入归一化到均值0方差1然后再做一次线性变换[ y \gamma \hat{x} \beta ]这里的gamma是每个通道一个可学习的缩放因子。如果某个通道的gamma值被训练到接近0那么这个通道的输出基本就是一个常数偏移对后续结果几乎没有贡献删除它是很合理的。基于这个思路网络瘦身Network Slimming方法会在训练损失函数里额外加入对gamma的L1正则惩罚让gamma尽量稀疏化。优化器在正常训练loss之外还会往gamma的反方向推把一部分通道逼到接近0。等训练结束后剪枝时只要统计gamma分布设置一个阈值把低于阈值的通道删掉就行。这套做法实现起来不复杂训练开销也小不需要额外的辅助网络或者复杂的显著性计算。对YOLOv8这种深层卷积网络特别适用。2.3 整体流程框架整个剪枝流程我总结为四个阶段稀疏化训练在正常训练目标上加入gamma的L1正则让BN缩放因子稀疏化通道重要度分析遍历模型所有带BN的层统计gamma值分布结构化剪枝根据剪枝比例生成每层保留通道的掩码重建模型结构微调恢复剪枝后的模型精度会有一定损失需要用原数据集短时间训练找回这四个阶段里最容易翻车的是第三阶段。因为YOLOv8里除了普通卷积还有C2f模块内的跨层连接和Detect检测头的分支结构这些地方的通道对齐关系必须小心处理。下面我逐个环节展开。3. 源码实现核心细节3.1 环境准备与依赖版本先说依赖。剪枝本身不需要太多额外库核心是PyTorch和Ultralytics。我测试的版本组合如下依赖项推荐版本备注Python3.8 / 3.103.8以上都行3.11有些算子编译麻烦PyTorch1.13 / 2.0.12.0以上对int8量化支持更好Ultralytics8.0.xx8.0和8.1的模型结构略有差异注意版本CUDA11.7 / 12.1按你的显卡驱动装onnx1.14导出阶段使用我踩过的第一个坑就是Ultralytics版本。8.0.40和8.1.x之间C2f模块的表示方式有微调如果用旧脚本加载新模型解析层名时会漏掉一部分。所以源码里建议锁定ultralytics8.0.127这个版本结构稳定网上资料也多。3.2 稀疏化训练给BN层加L1正则稀疏化训练是整个流程里最需要耐心的部分。核心代码不复杂关键是在每次反向传播后对gamma做额外惩罚。以ultralytics框架为例训练循环里可以拿到模型返回的loss然后再手动加上稀疏项import torch def sparse_loss(model, lambda_sparse1e-5): 计算BN层gamma的L1范数用于让gamma稀疏化 loss_sparse 0.0 for name, module in model.named_modules(): if isinstance(module, torch.nn.BatchNorm2d): # 只对包含gamma的BN层做惩罚 if module.weight is not None: loss_sparse torch.abs(module.weight).sum() return lambda_sparse * loss_sparse # 在训练循环中叠加 loss loss_box loss_cls loss_dfl loss sparse_loss(model, lambda_sparse1e-4) loss.backward() optimizer.step()这里lambda_sparse的取值很关键。太大会让精度掉很多太小gamma稀疏效果差。根据我的测试YOLOv8s在COCO类数据集上lambda_sparse取1e-4到3e-4比较合适。如果你自己数据集只有几个类别任务简单可以稍微加大到5e-4。有一个很常见的误区稀疏化训练不是只做几个epoch就够了。如果你是从头训练应该从第一个epoch就开始加稀疏惩罚训练到模型收敛。如果是加载预训练权重再稀疏化至少需要训练80到120个epoch否则gamma分布没有足够时间被惩罚项推开。训练完稀疏化模型后先跑一遍验证集记录此时的mAP。这个mAP是后面剪枝后微调效果的对照基准。3.3 通道剪枝与模型结构重建剪枝阶段我会遍历模型的所有卷积层找出每个卷积层对应的BN层然后根据gamma值排序来确定哪些通道要保留。YOLOv8的结构中有这些需要特别注意的地方C2f模块由多个Bottleneck组成Bottleneck里有shortcut分支剪枝时所有并行的卷积层输出通道必须保持一致Detect头有多个不同尺度的分支不能单独看每个分支的gamma否则会出现通道数不一致下采样层的通道数变化倍率需要和后续层匹配我写了一个简单的层解析函数把每层的依赖关系整理成字典def parse_model_channels(model): 返回每个卷积层对应的输入通道、输出通道以及依赖的层 channel_info [] prev_channels 3 # 输入图片RGB通道 for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): in_ch module.in_channels out_ch module.out_channels channel_info.append({ name: name, in_channels: in_ch, out_channels: out_ch, bn_attached: find_bn_after(model, name), # 找到紧接的BN层 }) return channel_info剪枝的核心是生成一个“保留通道掩码”。假设某个卷积层有128个通道我们想保留70%那就对它的BN层gamma按绝对值从大到小排序取前90个通道的索引作为保留集合。但有shortcut的地方不能简单这么算。对于C2f里存在残差连接的卷积组所有相关层必须共享同一个保留下标否则输出通道数对不上模型直接报错。为了处理这个问题我在源码里做了一步“通道对齐传播”记录每个卷积层的输入来源层如果某层的输入来自另一个剪枝后的层则它的输入通道保留下标由上游层决定对于残差相加的地方把两个分支的保留下标取并集举个例子一个Bottleneck模块的结构是x - conv1(1x1) - conv2(3x3) - 加上x - 输出如果conv1和conv2的输出通道都被剪枝最后相加时x的通道数必须等于conv2的输出通道数。所以x本身不需要剪枝但conv1的保留下标一定要由Shortcut分支的输入决定。这个逻辑在源码里是用一个依赖传播函数做的。剪枝时我不用PyTorch的register_buffer之类的东西而是直接重建全新模型。首选方法是用ultralytics的yaml配置思路生成一个剪枝后的新yaml然后从零加载预训练权重中对应保留的部分。具体做法是这样的剪枝后我会生成一个字典记录每一层的“原始权重索引到新模型索引”的映射然后拷贝权重def prune_state_dict(old_state_dict, keep_indices_map): new_state_dict {} for key, value in old_state_dict.items(): # 找到这个key对应的层和参数类型weight/bias/running_mean等 layer_name, param_type split_key(key) if layer_name in keep_indices_map: keep_idx keep_indices_map[layer_name] if param_type weight: # 输出通道对应BN缩放因子 if value.dim() 4: # 卷积权重 shape: [out_c, in_c, kh, kw] new_value value[keep_idx] elif value.dim() 1: # BN weight/bias等 new_value value[keep_idx] else: new_value value # 其他情况 else: new_value value new_state_dict[key] new_value else: # 没有被剪枝的层原样复制 new_state_dict[key] value return new_state_dict这里有一个关键点如果剪掉了输入通道那么当前卷积层的输入权重也要按上一个层的保留下标来索引。也就是卷积权重要同时做两次索引一次是输出通道一次是输入通道。很容易写错。完整的权重映射函数如下def remap_conv_weight(weight, keep_out_idx, keep_in_idx): weight shape: [out_c, in_c, kh, kw] 按输出通道保留下标和输入通道保留下标分别筛选 return weight[keep_out_idx][:, keep_in_idx, :, :]注意对于1x1卷积和3x3卷积都一样处理因为kh、kw维度不受影响。3.4 微调恢复精度剪枝完成后千万不要直接拿剪枝后的模型任何数据都不跑就部署。因为gamma稀疏化训练会把一些重要通道的值也压小直接剪掉会导致特征表示不完整。必须做微调。微调的epoch不需要像从头训练那么长一般用原学习率的1/10训练20到30个epoch即可。比如原训练lr是0.001微调用0.0001。权重衰减可以不变但稀疏惩罚项在微调阶段要去掉否则又会让gamma稀疏影响恢复。微调时我习惯把模型结构保存为新的yaml并用ultralytics的DetectionTrainer加载权重。此时需要注意剪枝后的模型权重键名要和yaml结构的层名一一对应。如果出现加载错误多半是yaml中层的定义和权重shape对不上。微调完再跑验证集正常情况下mAP能恢复到剪枝前mAP的97%以上。如果差得比较多检查剪枝比例是否太大或者微调lr是不是没调对。4. 实操过程与参数调优4.1 数据准备与训练超参数设置YOLOv8剪枝训练和你平时的训练流程差不多数据集格式还是一样的。我这里给一个具体的配置参考以你自己的数据集为例# dataset.yaml train: /path/to/train/images val: /path/to/val/images nc: 5 names: [person, car, bike, dog, cat]稀疏化训练的超参数我用的是epochs: 100从预训练权重开始的话batch_size: 32如果显存有限就16但BN统计量会受影响lr0: 0.001weight_decay: 0.0005lambda_sparse: 2e-4优化器: SGD比Adam更好控制稀疏化gamma分布更干净SGD配合L1正则的效果比Adam稳定。Adam对每个参数单独调整学习率会削弱L1惩罚的力度稀疏化效果不理想。我刚开始用Adam试过gamma分布拖泥带水阈值不好定换成SGD之后顺滑很多。4.2 剪枝比例选择策略不要一刀切这是整个项目里最需要经验和耐心的环节。很多教程会告诉你“剪去30%”或者“剪去50%”但实际不是每个层都剪相同比例才最优。深度学习网络有个特点靠近输入层的特征表征往往更通用靠近输出层的特征更任务专属。对不同层采用统一比例剪枝会导致浅层表达能力严重受损而深层可能还有冗余。我推荐的做法是采用分段比例骨干网络的前几层剪枝比例控制在20%以内骨干网络的深层剪枝比例可以到40%到50%C2f模块内部保持30%左右Detect检测头尽量少剪最多剪10%到15%检测头虽然参数量占比不大但直接影响最终输出的类别和回归精度剪多了mAP很难回来。这也是我踩过最深刻的坑第一次统一剪50%时检测头直接给剪残了mAP掉了10多个点。为了自动确定每层的比例我在源码里加了一个简单的策略统计每层BN的gamma绝对值总和按照总和的反比例分配剪枝强度。总和小的层说明本身重要度集中少剪总和大的层说明冗余度高多剪。4.3 针对YOLOv8结构特点的定制处理YOLOv8的Detect头和上一代YOLOv5有区别。YOLOv8的检测头采用了解耦结构分类和回归分支分开同时还用DFLDistribution Focal Loss。在这个框架下检测头里包含很多卷积层通道关系比v5复杂。剪枝检测头的时候必须把分类分支和回归分支对应的输出通道做一致性处理。因为最终预测时每个分支的输出通道数会映射到类别数和锚框相关的参数上如果剪完通道数对不上后续解码逻辑就崩了。源码中对检测头统一做了“只修剪内部隐藏层不修剪最终输出层”的处理。另外C2f模块里有一个cv1卷积、n个Bottleneck、一个cv2卷积。Bottleneck里的两个卷积层输出通道必须一致否则残差相加没法做。而cv2卷积的输入通道是所有Bottleneck输出通道的拼接。剪Bottleneck时必须保证所有Bottleneck用同一个保留下标集合。这个我在前面解析依赖时已经强调过这里再加深一下印象。我实际操作中发现一个高效方案直接修改ultralytics的yaml把C2f模块中的Bottleneck数量改少同时把剪枝后的通道数填到yaml里这样模型结构更干净。但这样改出来的结构无法复用原来的权重需要从零训练耗时比较长。所以我的源码里还是保留权重拷贝方式能节省时间。4.4 剪枝后模型的导出与验证剪枝完成后除了验证mAP我还建议做两件事导出ONNX确认计算图正确用目标推理框架跑一次延迟。导出ONNX时有几个参数值得注意model.eval() dummy_input torch.randn(1, 3, 640, 640).to(device) torch.onnx.export( model, dummy_input, pruned_yolov8s.onnx, opset_version12, input_names[images], output_names[output0], dynamic_axes{ images: {0: batch}, output0: {0: batch} } )如果你用静态batch部署建议把dynamic_axes去掉这样TensorRT优化时能做更多层融合。如果要用动态batch才保留dynamic_axes。导出之后我自己会用onnxruntime跑一遍python -c import onnxruntime as ort import numpy as np sess ort.InferenceSession(pruned_yolov8s.onnx, providers[CPUExecutionProvider]) x np.random.randn(1,3,640,640).astype(np.float32) y sess.run(None, {images: x}) print(输出shape:, [o.shape for o in y]) 如果输出shape正常就说明剪枝后的模型在计算图上没有结构问题。再进一步就是转成TensorRT的engine文件用trtexec测试延迟。TensorRT对剪枝后的模型通常很友好因为通道数减少后每层计算量降低能获得比PyTorch环境更好的加速。5. 常见问题与排查技巧实录5.1 BN层gamma不稀疏剪了等于白剪这是最常遇到的问题。稀疏化训练结束后可视化gamma直方图发现都集中在0附近但不是真正接近0或者分布很均匀。这说明L1正则的力度不够或者训练epoch太短。解决办法检查代码里是否正确遍历到了所有BatchNorm2d层。YOLOv8内部有些模块可能用了可重参数化的结构部分BN层不在列表里。增大lambda_sparse但每次只增大1.5倍左右不要一下翻十倍。训练epoch拉长我遇到过一个数据集从80epoch加到150epoch后gamma稀疏度明显改善。真正适合剪枝的gamma分布应该像梳子一样大部分接近0少部分明显大于0。如果看不到这种分布不要强行剪枝。5.2 剪枝后mAP掉得厉害剪枝掉几个点是正常的但如果掉了超过5个点就需要反思几个因素。第一个是剪枝比例是否过高尤其是检测头和C2f模块。前面说过检测头少剪。第二个是微调epoch和lr是否合适。我发现微调lr过大反而会破坏已经学到的特征导致精度无法恢复。微调阶段应该让模型在原权重附近小范围移动不要跑远。第三个是训练稀疏化时的batch size。BN的统计量在稀疏化训练时很重要batch size太小会引入噪声让gamma值看起来该剪的没剪、不该剪的却很小。我用batch size 32比用16的效果好有条件可以用64。5.3 剪枝后模型结构输出shape对不上这类错误通常发生在C2f的Bottleneck和Detect分支。如果只按单个卷积层的gamma排序剪枝没有考虑残差连接那最终拼接通道数必然出错。我的排查方法是剪枝后先用一个小测试输入打印每一层输出的shape逐层比对。定位到第一个shape不匹配的位置再回头检查依赖关系图。源码里专门写了auto_check_shape这个函数def auto_check_shape(model, input_size(1, 3, 640, 640)): x torch.randn(*input_size) hooks [] def hook_fn(module, input, output): print(f{module.__class__.__name__}: {output.shape}) for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.BatchNorm2d): hooks.append(module.register_forward_hook(hook_fn)) model(x) for hook in hooks: hook.remove()这个脚本能帮你快速定位是哪一层的通道数出了问题不用在yaml里猜。5.4 剪枝后推理速度没有明显提升这个问题解决思路要分两步走。第一步确认剪掉的FLOPs是真的从计算图中消失了。有些工具虽然报告剪枝比例但PyTorch模型由于结构里有Fan-in、cat操作实际计算路径中可能有额外的零填充或者换位操作导致前向时间没降。可以重新导出一遍ONNX用netron看看计算图结构。第二步考虑硬件对齐。GPU推理对通道数通常以32或64为单位比如TensorRT的某些kernel会按8的倍数对齐。如果你的剪枝保留了比如37个通道反而比64个通道更慢。所以剪枝时可以做一个“通道圆整”处理把保留下标数量对齐到8或16的整数倍。源码中支持设置align_channels16按这个参数自动调整剪枝策略。我做过一个案例不圆整剪到30%时FLOPs降了很多但TensorRT延迟几乎没变把通道圆整到16的倍数后延迟才真正降下来。这个属于实战中很容易忽略的点值得记住。5.5 预剪枝和后剪枝的区别我看到网上也有人在问“预剪枝与后剪枝”。这两个概念在模型压缩里指代不一样但很多教程把“稀疏化训练-剪枝-微调”称为训练后剪枝而“从头训练一个小网络”称为预定义剪枝。严格意义上说传统决策树的预剪枝和后剪枝不是同一套逻辑但在神经网络压缩领域大家通常关心的是剪枝时机。我的做法始终是预训练权重 - 稀疏化训练 - 剪枝 - 微调。这套流程本质上是训练后剪枝。如果你想省时间直接在预训练权重上稀疏化训练几十个epoch也可以但效果会略差一点。预剪枝先定义小结构再训练除非你有足够算力从头训练否则效果不一定好因为YOLOv8很多层之间有复杂的耦合直接定小结构容易欠拟合。6. 核心源码结构与扩展方向6.1 源码文件组织我把这套剪枝源码整理成了几个模块结构大概是这样的yolov8_prune/ ├── sparse_train.py # 稀疏化训练入口集成ultralytics训练循环 ├── prune_analyze.py # 统计BN层gamma分布生成剪枝报告 ├── prune_channels.py # 通道剪枝与权重重映射 ├── rebuild_model.py # 重建剪枝后的model和yaml ├── finetune.py # 微调剪枝后模型 ├── export_onnx.py # 导出ONNX └── utils/ ├── dependency.py # 解析模型层依赖关系 ├── remap.py # 权重索引映射 └── shape_check.py # 自动shape检查这个结构是按流程拆的每一步单独执行方便排查。不建议把稀疏化和剪枝写在一个脚本里否则出了问题不好定位。6.2 配合量化的扩展建议剪枝做完之后通常接着做量化两者配合能进一步压缩模型。剪枝减少通道数量化降低每个参数的比特数二者正交。比如yolov8s剪枝35%后再做INT8量化模型体积能压到原始体积的10%左右。不过要注意剪枝和量化叠加会出现精度下降叠加。所以量化一般要做校准拿几百张代表性图片收集激活值范围。剪枝后的模型重新量化时因为我之前提过BN折叠的问题最好先把BN融入卷积再量化否则量化误差会很大。6.3 扩展到大模型和更多任务的思路这套源码虽然针对yolov8s但核心逻辑可以复用在yolov8m、yolov8l甚至YOLOv5上。关键点还是依赖关系解析和保留下标传播。我自己试过将剪枝逻辑迁移到yolov8m只需要修改yaml解析的映射表流程不变。如果你以后要跑更强的检测网络比如YOLOv9或YOLOv10思路也一样但需要重新对照网络结构里的跨层连接方式。每个模型的残差和拼接方式不同剪枝脚本没法做到“一劳永逸”这也是为什么做模型压缩的人必须看得懂网络结构不能只靠工具包。我个人在实际操作中体会最深的一步其实是剪枝完成后的模型验证。很多人急着看加速效果但我会先在验证集上跑完mAP再导出ONNX测试完整性最后才上推理框架测延迟。这三步按顺序走完剪枝后的模型才敢上线。这套流程看起来繁琐但能省下后面排查问题的大量时间。如果你正在准备剪自己的检测模型我建议先拿一个小数据集跑通全流程再换真实数据调参。毕竟剪枝这活儿参数选不好代码再对也出不来效果。
返回列表