★ FEATURED ARTICLE 模型压缩这件事我一直觉得是YOLO系列落地到边缘设备时最值得系统整理的一环。前段时间帮一个嵌入式项目做检测模型迁移模型在服务器上测试FPS接近七十换到目标板卡一跑直接掉到个位数CPU占用拉满推理一帧要等小几百毫秒。裁输入分辨率、换推理引擎、做动态量化都试了一遍真正把速度救回来的还是模型剪枝和知识蒸馏这一套组合拳。这篇文章以YOLO系列为具体对象把剪枝和蒸馏的原理、落地步骤、超参调节和踩坑记录完整写出来给打算做模型压缩部署的朋友一份可以直接照抄的作业。1. 为什么YOLO在端侧跑不动剪枝和蒸馏各自解决什么问题YOLO系列本身已经是效率极高的单阶段检测器但效率和“能在端侧实时跑”是两码事。以一个参数量大约七百万的轻量YOLO为例推理一次的FLOPs仍然可能高达十几个G而这只是计算量真正到了端侧算力、内存带宽、缓存命中率、算子支持度每一项都会拖后腿。很多边缘芯片的整数算力可能不错但浮点算力只有服务器的几十分之一模型推理从几十毫秒变成几百毫秒一点也不奇怪。很多人第一反应是“换一个小模型”比如从大型版本换成微型版本或者从零训练一个窄网络。这个思路没错但往往踩到一个坑从零训练的小模型在同等数据量和训练周期下精度通常达不到大模型压缩后的水平。原因不复杂——大模型的参数里已经包含了对数据分布更完整的拟合这些信息很难靠随机初始化的小模型自己重新“悟”出来。这时候就需要模型压缩技术而剪枝和蒸馏就是两条最常用的路线。用一句话概括两者的差别剪枝是做减法把已经训练好的大模型里冗余的通道、权重删掉留下一个更瘦但功能几乎不变的结构蒸馏是做转移让一个小模型去模仿大模型的行为把大模型的“经验”迁移到小模型身上。两者可以单独用但YOLO落地上最常用的方式是把它们串起来——先剪掉冗余结构再由大模型当老师把精度“喂”回来。我见过不少团队把这两件事混为一谈或者只做其中一步就觉得万事大吉。实际上它们解决的问题并不重叠剪枝解决的是结构冗余蒸馏解决的是参数信息不足。给一张对比表应该能把这套逻辑说清楚。方案对结构的影响精度恢复手段适用场景从零训练小模型直接使用紧凑结构依赖训练技巧与大算力时间充足、任务简单大模型通道剪枝删除冗余通道剪后微调结构明显冗余、需要激进压缩大模型蒸馏小模型结构不变用大模型软标签监督小模型参数有限但想要高精度实际项目里剪枝之后往往伴随精度掉点微调能补回来一部分但补不干净蒸馏则可以在不改变结构的前提下进一步提升小模型的上限。所以真正工程化的路线是先剪枝把模型变小再用原版大模型做蒸馏去补救精度损失。后面几章就围绕这条路线展开。2. 通道剪枝在YOLO上的落地从BN稀疏化到依赖层分组2.1 BN层的gamma系数为什么是通道重要性的最佳度量通道剪枝的关键问题只有一个怎么判断哪些通道是冗余的。目前工业界最稳定的做法是用BN层的gamma系数作为通道重要性的衡量指标。YOLO系列每个卷积层后面几乎都跟着BN层而BN层在做归一化之后会做一次缩放和平移y gamma * x_hat beta其中x_hat是归一化后的特征gamma是每个通道的缩放系数。如果某个通道的gamma趋近于零那不管输入是什么这个通道的输出都接近常数beta对后续特征几乎没有贡献。这类通道就是典型的冗余通道删掉它们对整体精度影响很小。为什么不用卷积核的L2范数来判断因为卷积核的参数分布受上一层输出尺度影响不同层的卷积核范数不在同一个量级上跨层比较没有意义。而BN的gamma直接作用在归一化后的特征上天然处于同一尺度可以直接拿来排序。这个细节很重要很多人一开始用卷积核范数剪枝结果不同层的保留比例差距极大精度崩得很厉害。2.2 稀疏化训练把冗余通道逼到墙角直接从预训练模型上按gamma排序剪枝效果通常不好。原因是预训练模型的gamma分布比较均匀很难找到一条清晰的“分界线”。你设定剪枝比例后被剪掉的那部分通道仍有一定的贡献剪完必然掉点。解决办法是稀疏化训练在原有检测损失上对BN gamma施加L1正则惩罚让大部分gamma在训练过程中趋近于零。L1范数优化的特点就是产生稀疏解而不是像L2那样把所有值均匀地缩小。直观理解就是训练时每条通道都有一个“使用成本”模型会自动放弃那些不太重要的通道。稀疏化训练的损失函数大概长这样# 在已有检测loss基础上叠加L1稀疏惩罚 bn_gammas [m.weight for m in model.modules() if isinstance(m, nn.BatchNorm2d)] sparse_term sum(torch.abs(g).sum() for g in bn_gammas) loss det_loss sparse_ratio * sparse_term稀疏系数sparse_ratio不建议一开始就拉满。我习惯从0.0005左右开始训练二三十个epoch后观察gamma分布直方图如果大部分gamma仍然集中在1附近再适当加大系数。理想状态是gamma分布呈现明显的两极分化一部分接近0另一部分保持在正常范围。这个过程一般需要五十到一百个epoch直接使用预训练模型继续训练即可不需要重新从零开始。2.3 shortcut与concat带来的索引依赖YOLO剪枝最关键的坑YOLO系列的结构远不是简单的卷积堆叠里面布满了shortcut残差连接和concat拼接操作。残差结构要求两个分支的通道数完全一致PAN结构的特征融合则把多个分支的特征在通道维度直接拼接。这就带来一个剪枝核心问题不能单独决定某一层的通道保留情况必须保证所有关联层使用同一套保留索引。举个例子某个C3/C2f模块里的瓶颈结构主分支经过卷积后与shortcut分支相加。如果你按gamma排序把主分支某个卷积的通道剪掉一批shortcut分支的通道数没有同步变化两个分支就无法对齐模型直接报维度错误。concat操作也一样拼接后的通道数必须与后续卷积的输入通道保持一致。所以正规的YOLO剪枝流程要先做“依赖层分组”。简单说就是反向遍历网络结构找出所有直接或间接共享张量维度的卷积层把它们归为一组。剪枝时以组为单位合并组内所有BN的gamma统计量再统一排序每一层都套用同一个保留索引。分组这一步是YOLO剪枝和普通CNN剪枝最大的区别也是最容易翻车的地方。# 伪代码构建“同一裁剪索引”的层分组 def build_prune_groups(model): groups [] # 找到所有Concat和shortcut分支记录它们对应的BN层编号 # 把输入来自同一concat/shortcut的卷积层的BN index放入同一group for node in reversed(model_graph): if node.type in (concat, add): new_group collect_dependent_bn(node.inputs) groups.append(new_group) # 每个group内部合并gamma统计量排序后生成统一保留索引 return groups这一块必须写在脚本里自动完成不要人工判断。YOLO的neck部分层数多、拼接关系复杂人工数很容易漏。2.4 剪枝后的模型重建与精度恢复完成通道筛选之后需要按照保留索引重建一个更窄的模型。这里有两种实现方式一种是直接修改网络定义文件把每层的输出通道数改掉另一种是动态地从原模型的state_dict里按索引抽取参数生成新模型。推荐第二种因为自动化和可复现性更好。重建时最关键的操作是参数抽取卷积的权重按输入通道和输出通道两个维度分别筛选BN的gamma、beta、running_mean、running_var都按输出通道筛选。以torch为例本质上就是一行index_select的事但顺序绝对不能错。比如卷积层的weight shape是[out_channels, in_channels, k, k]先按保留索引过滤out维度再按上一层传来的in维度过滤输入。剪完之后一定不要直接拿去部署精度大概率是掉的。需要先做一轮“恢复性微调”学习率比正常训练小一个数量级训练二十到三十个epoch。如果稀疏化训练做得充分微调后精度通常能回到原来的98%以上。之后如果还想进一步压榨再接入蒸馏阶段。def rebuild_model_by_mask(old_model, keep_masks): new_model create_scaled_model(old_model, keep_masks) for (name, old_param), (new_name, new_param) in zip( old_model.named_parameters(), new_model.named_parameters() ): if weight in name and conv in name: new_param.data old_param.data[keep_masks[name]].clone()3. 检测任务的蒸馏为什么不能只照抄分类结果3.1 分类蒸馏的基本逻辑知识蒸馏最早在图像分类任务上被验证有效。核心思想是大模型的输出中除了正确答案之外那些“错误类别上的概率分布”同样包含了信息。比如一张图被分类为猫大模型可能给出“猫0.7、狗0.2、老虎0.1”这些概率反映出“狗和老虎在特征上接近猫”。让一个小模型去学这种软化的概率分布比只学硬标签能获得更多监督信号。具体做法是引入温度参数T先对logits做softmax软化再计算学生和老师分布的KL散度def kd_loss(student_logits, teacher_logits, T): student_soft F.log_softmax(student_logits / T, dim-1) teacher_soft F.softmax(teacher_logits / T, dim-1) return F.kl_div(student_soft, teacher_soft, reductionbatchmean) * T * TT越大概率分布越平滑小模型越能学到类别之间的相似关系。公式最后乘上T的平方是为了抵消softmax求导时引入的1/T因子让梯度尺度保持稳定。这个细节很关键我在项目里见过有人忘记乘T^2结果蒸馏loss在数值上被压得太小完全不起作用。3.2 YOLO蒸馏的三个落点分类logits、neck特征、解码头分类蒸馏的直接套用在YOLO上行不通。检测任务的输出由类别概率、边界框坐标、目标置信度三部分组成其中坐标回归是连续值根本不存在类别分布而一张图上正负样本比例可能达到几百比一大量负样本区域的logits会淹没真实目标的信息。所以YOLO的蒸馏必须拆开设计分三个位置落点。第一是分类分支的logits蒸馏这个和分类任务最接近但只在检测头认为的“正样本位置”上计算负样本位置要么忽略要么用很低的权重。第二是neck层特征的蒸馏把teacher和student在P3/P4/P5三个尺度的输出特征做空间对齐用MSE loss拉近。由于teacher的特征通道数通常比student多一般要先加一个1x1卷积把teacher投影到student的通道维度而不是反过来硬截断。第三是解码头的box回归蒸馏让student的预测框坐标逼近teacher的坐标只在正样本上计算L1或GIoU lossteacher的目标置信度可以作为权重。用一个表格快速总结蒸馏位置对齐对象损失函数备注分类logits类别概率分布KL散度只在正样本位置生效Neck特征P3/P4/P5特征图MSE / L2需要1x1卷积投影对齐通道解码头坐标与置信度L1 / GIoU BCE只对正样本计算最终的训练总损失可以写成total_loss det_loss alpha * feat_distill_loss beta * logits_kd_loss gamma * box_distill_loss三个权重不是拍脑袋定的。我在实践中会把feat_distill_loss的权重放在0.1到0.5之间logits_kd_loss的权重小一些box_distill_loss的权重再小一个量级。如果特征蒸馏权重太大student会过度模仿特征的绝对值反而忽略了自己检测头的定位优化框的位置容易偏移。3.3 Teacher模型的用法与容易忽略的细节蒸馏训练时teacher必须冻结并且要切到eval模式。这里有个很多人忽略的细节如果teacher里还有Dropout或者带随机性的模块必须用eval模式关闭否则每次前向传播的软标签都不一样student会学到一个“抖动”的目标。更稳妥的做法是使用teacher的EMA权重版本也就是训练过程中指数滑动平均产生的权重通常比最终权重更平滑蒸馏出来的学生精度更高。另外teacher和student的输入预处理必须完全一致。YOLO训练时常用的mosaic、mixup增强会让teacher看到拼接后的图像导致soft label的分布和质量发生变化。如果要做蒸馏建议在启用mosaic的batch上关闭蒸馏loss或者在数据加载阶段对teacher使用和student完全一样的增强参数。加大蒸馏温度T也能缓解这个问题但根子还是数据一致性。4. 剪枝与蒸馏怎么组合才能把掉点救回来4.1 先剪枝后蒸馏最实用的组合路线剪枝和蒸馏前面分别讲清楚了但工程上的核心问题是顺序。我自己的推荐是先稀疏化剪枝再用原版大模型当teacher做蒸馏微调。整个流程四步准备一个训练充分的大模型既当剪枝源又当后续的teacher。对模型做稀疏化训练把BN gamma拉开。按剪枝比例完成通道裁剪和模型重建。用未剪枝的原模型冻结做teacher对剪枝后的student做检测蒸馏微调。为什么这个顺序最实用因为剪枝会让模型掉点而蒸馏恰好是修补掉点最有效的手段。teacher和student来自同一个初始模型结构上高度相似蒸馏目标的“迁移距离”短学生学起来比从随机初始化的小模型容易得多。我在某次安防小目标检测实验里做过对比剪掉一半FLOPs之后只做普通微调mAP掉了约4个点同样剪枝比例在微调阶段加入蒸馏掉点控制在1.5个点以内。这个差距在真实项目里非常明显。4.2 另外两条路线先蒸馏后剪枝和迭代式剪枝也有人用“先蒸馏后剪枝”的顺序先让一个小模型从大模型那里把精度学到手再对这个已经紧凑的模型做剪枝。这个路线适合目标模型一开始就非常小的情况但流程更长中间产物多剪枝阶段掉点后没有更大的teacher可以依赖。迭代式剪枝是另一种思路先剪10%微调恢复再剪10%再微调恢复每一步的剪枝比例都很小所以单步精度的损失可控。效果确实好但训练周期可能拉长到两到三倍。如果项目周期不紧可以试试如果追求一次到位我建议一次性剪到目标比例配合蒸馏微调性价比最高。4.3 调参优先级先定剪枝率再定蒸馏权重这套组合拳里面有大量超参稀疏化系数、剪枝比例、蒸馏温度、特征蒸馏权重、box蒸馏权重。如果同时乱调最后出了问题根本不知道是哪个环节造成的。我建议按下面的优先级顺序来调先确定剪枝比例这是由目标FLOPs和延迟决定的属于硬性约束。再调稀疏化系数保证gamma分布拉开的同时精度不掉太多。最后调蒸馏权重先用一组保守值跑通再按mAP曲线调整。蒸馏权重也不要从头到尾固定不变。我常用的一个技巧是前十个epoch只做普通微调后二十个epoch再逐步打开蒸馏loss让模型先稳定自身结构再跟着teacher走。这样比全程开启蒸馏更稳定。5. 踩坑清单模型变“瘦”之后才暴露出来的问题5.1 剪枝后精度崩盘先查BN统计量有一次剪完模型后mAP直接掉了十几个点比预计严重得多。排查到最后问题出在BN上。稀疏化训练中L1正则会把gamma往零推这个过程中BN的running_mean和running_var也会跟着变化。剪枝重建后新模型里这些统计量是从旧模型直接抄过来的但它们对应的通道已经被换掉了分布完全对不上。解决办法是在剪枝重建后先用少量训练数据重新前向走一遍只更新BN统计量、不更新梯度让新的mean和var适配新的通道结构。这一步有时比微调本身还重要。5.2 shape mismatchshortcut组没对齐这个是YOLO剪枝最常见的报错RuntimeError提示某个cat或add操作维度不一致。原因几乎都是依赖层分组写漏了只剪了主干分支的卷积没把shortcut分支的卷积纳入同一组。排查方法是打印每个concat/add节点的输入shape找到是哪一层出现的分叉然后把对应的BN层index补进同一个prune group。遇到这种情况就写一个更全面的依赖分析脚本把同组识别做成自动化不要靠人肉看结构图。5.3 Detect head被剪出“畸形”输出YOLO的检测头输出通道数和类别数、锚框数强绑定解码逻辑对输出张量的shape有硬性要求。如果你对检测头里的卷积也做了通道剪枝输出shape就变了轻则训练报错重则模型能训练但推理阶段解码彻底乱掉。最稳妥的方案是剪枝范围只覆盖backbone和neck部分Detect head完整保留。检测头本身节点很少剪它对总计算量帮助有限但破坏输出的可能性极高。只要剪掉backbone侧的输入通道处理好neck到head之间的衔接层即可。5.4 蒸馏权重过大检测框开始漂移某次实验里加了特征蒸馏后分类准确率明显变好但mAP反而下降了。观察预测结果发现框的位置偏移变大了。原因是特征蒸馏鼓励student的特征图向teacher靠拢但并没有约束这些特征解码出来的框一定要准。特征蒸馏和检测任务之间存在“目标竞争”。解决办法是把box回归的蒸馏权重提上来或者在训练后期衰减特征蒸馏的权重把主导权还给检测loss。5.5 mosaic增强让teacher的软标签不稳定如果蒸馏loss曲线一直上下乱跳可能是mosaic增强在捣乱。mosaic会在一张图里拼接四张不同来源的图片teacher看到的目标分布和学生看到的不一样特别是小目标被拼接边界切割后teacher输出的置信度不可靠。我当时的处理是在数据加载器里加一个开关如果当前batch启用了mosaic这一批就不计算蒸馏loss只算普通检测loss。同时把mixup概率从0.5降到0.1软标签的稳定性立刻好很多。5.6 稀疏化系数过大模型变“植物人”L1稀疏强度不是越大越好。有次我把稀疏系数调到0.01训练几十轮后检查gamma分布发现几乎所有通道的gamma都被压到接近零模型输出直接变成一个常数怎么微调都救不回来。稀疏化训练过程中要定期看gamma直方图如果发现绝大多数gamma聚集在0附近说明稀疏过头了必须回退到更小的系数重新训练。理想状态是大概一半以上通道保持正常值一小部分被压到接近零。6. 一套可以直接复制的YOLO压缩流程模板说了这么多最后给一套可以直接用的流程模板涵盖稀疏化、剪枝、蒸馏三个阶段具体超参可以按项目调整。第一阶段稀疏化训练。加载预训练权重沿用原训练配置在loss中加入BN gamma的L1正则系数从0.0005开始训练50到100个epoch每隔10个epoch打印一次gamma分布直方图目标gamma出现明显两极分化且mAP不显著下降第二阶段通道剪枝。扫描模型结构构建依赖层分组统一统计每组gamma按保留比例生成每层的keep mask重建模型按mask抽取权重和BN参数用少量数据重估BN统计量再开始微调第三阶段蒸馏微调。加载原版大模型作为teacher冻结并切换eval模式student加载剪枝后的模型先做10个epoch普通微调逐步加入分类logits蒸馏、neck特征蒸馏、box回归蒸馏最终模型验证FLOPs、参数量、mAP和端侧延迟下面是推荐参数区间标好调整方向参数推荐取值调整方向稀疏化系数0.0005 ~ 0.001过大则gamma整体趋零过小则分布不拉开剪枝比例20% ~ 50%超过50%需要非常充分的蒸馏配合蒸馏温度T4 ~ 6T过小近似硬标签T过大会丢失细节特征蒸馏权重0.1 ~ 0.5权重过大会导致检测框偏移box蒸馏权重0.05 ~ 0.2只对正样本计算微调epoch数30 ~ 50前10轮过渡后20轮稳定最后说一个个人习惯在稀疏化训练阶段我会先把稀疏系数设得很小比如0.0002然后观察一轮mAP掉点是否在0.3以内再慢慢往上加。不要直接上高系数否则万一训练已经跑了几十个epoch才发现模型被压坏回退成本很高。剪枝比例也建议做成可配置参数先剪一个保守比例跑通整个流程确认链路没问题后再逐步加大。模型压缩这件事最怕的不是剪得不够多而是流程跑不通的时候你根本分不清是哪个环节出了问题。
阅读完成 · 觉得有帮助？
