模型压缩与推理加速：剪枝、量化、蒸馏一体化实践

算起来我在组里的角色一直有点“擦屁股”的意思。算法同事训完模型扔给我一句“效果不错”然后模型就归我管了。一开始我以为管的就是部署那点事后来才发现真正让人头疼的不是模型跑不起来而是跑起来了但太慢、太大、太贵。就在这个背景下我动手折腾了一个叫“Model-Optimizer”的内部小工具集。如果你也在搞模型压缩、推理加速这类事儿或者正在为了一个BERT模型的内存占用掉头发这个分享应该能给你一些省时间的思路。Model-Optimizer说白了就是一套把“训练好的模型”变成“能高效部署的模型”的流程化工具。它可以对模型做结构性剪枝、动态/静态量化、知识蒸馏的召回以及一键导出ONNX格式做推理验证。它不是某个单一算法更像是一个管线把业内常说的剪枝、量化、蒸馏都串了起来。适合谁适合那些已经能训练出模型但卡在部署环节的工程师也适合想在模型体积和推理延迟之间找平衡点的同学。下面我把整个设计思路、实操过程和踩过的坑都摊开说。1. 为什么我会写一个Model-Optimizer1.1 第一次被线上推理延迟逼疯的经历故事得从一次线上事故讲起。当时组里的OCR服务用的还是ResNet50作为特征提取骨干单张图在GPU上的推理时延大概在12ms到15ms之间CPU上则直接飞到200ms以上。换成生产环境的多租户混合部署后单卡吞吐根本扛不住流量峰值机器成本一个月涨了小十万。老板拍板要么换更贵的卡要么把模型瘦身。那次我尝试了最简单的办法——直接用PyTorch的torch.jit.script转成TorchScript结果只优化了不到10%的延迟因为模型本身的红利早就被框架优化过了。于是我开始认真调研剪枝和量化方案。后来陆陆续续试过[NNI]、[Intel的蒸馏工具]、[ONNX Runtime]以及一些开源库但问题在于这些工具各自为政剪枝是剪枝量化是量化导出是导出没有一条顺畅的流水线。每次都要写一堆黏合代码而且换了模型结构之后前面的配置全得推翻重来。这也催生了Model-Optimizer的第一版一个用YAML配置驱动、模块化调用剪枝/量化/蒸馏/导出的工具集。起初它只是我个人的脚本仓库后来组里其他人也开始用逐渐加上了自动评估和回归对比的能力。1.2 模型优化到底在优化什么在继续聊工具之前得先把一个问题说透:所谓“模型优化”到底优化的是什么我把大方向的优化目标拆成了四类:体积模型文件多大内存占用多少。主要是参数量和位宽问题。速度单次推理的延迟Latency或者单位时间能处理的样本数Throughput。受算子的计算量、访存量、并行度影响。能耗在边缘设备、手机端比较敏感模型越重功耗越高。精度保持压缩之后不能掉点太多。通常我们会设一个可接受阈值比如精度下降控制在1%以内。Model-Optimizer的核心逻辑就是在这四个维度之间找平衡。剪枝负责砍掉冗余参数量化负责把FP32变成FP16或INT8知识蒸馏用大模型带小模型ONNX Runtime负责把图优化和算子融合做干净。四者叠加在一起体积最多能压到原来的1/4以下延迟在CPU上能砍掉50%到70%。需要强调一下:模型优化和训练时的“调参优化”完全不是一回事。调参是在优化损失函数让模型收敛得更好模型优化是在模型已经训好的前提下做推理侧的工程优化。前者是炼丹后者是锻钢不少人一开始把这两个概念混在一起就会走弯路。2. Model-Optimizer整体设计与方案选型2.1 优化手段的组合剪枝、量化、蒸馏各司其职Model-Optimizer把整个优化管线分成了四个核心模块:结构化剪枝、量化压缩、知识蒸馏和格式导出。一开始我也想过是不是只做其中一部分就行但实际组合之后发现效果是乘法的关系不是加法。剪枝我重点关注的是结构化剪枝。以卷积网络为例BN层BatchNorm里的缩放因子gamma可以用来衡量每个通道的重要性。如果某个通道的gamma值普遍很小说明它对最终输出的影响有限剪掉它对精度的影响相对可控。这个思路参考了“Learning Efficient Convolutional Networks through Network Slimming”那篇论文的思想。在Model-Optimizer里我实现了对Conv2d和BatchNorm2d的通道级别剪枝可以直接把卷积层的输出通道数缩小属于真正的“瘦身”。量化这块Model-Optimizer提供两种路径Post-Training QuantizationPTQ和Quantization-Aware TrainingQAT。PTQ适合快速验证直接把训练好的FP32模型转成INT8不需要重训但精度损失靠校准集来缓解。QAT则是在训练过程中模拟量化误差精度保持更好但需要重新走一遍训练时间成本高。工具内部约定好了规则小模型、强结构模型优先PTQ大模型或者精度敏感场景选QAT。知识蒸馏模块不是Model-Optimizer首创的东西但我把它和剪枝量化放在了一条链路上:先用大模型蒸馏出一个较小的模型再做剪枝最后做量化。蒸馏我用的是软标签蒸馏让student模型去拟合teacher模型输出的概率分布而不仅仅是hard label的交叉熵。这样小模型学到的不光是“是什么”还有“像什么”在剪枝和量化之后抗掉点能力更强。2.2 为什么选择PyTorch ONNX Runtime这套组合工具选型上我几乎没有纠结过PyTorch生态在模型导出和动态图调试上的优势太明显了。训练、验证、微调都能在一套代码里完成TorchScript和ONNX导出的支持也都相当成熟。Model-Optimizer内部的主体代码用PyTorch最终产物面向推理时统一导出到ONNX格式。ONNX Runtime则是推理侧的“最后一公里”。它在CPU、GPU和NPU上都有对应的执行后端自带层级图优化比如算子融合ConvBNReLU融合成ConvFused、常量折叠、冗余节点消除等。光是把PyTorch模型导出成ONNX再喂给ONNX Runtime通常就能白捡10%到20%的加速这还不算量化和剪枝带来的收益。为什么不用TensorRTTensorRT在NVIDIA GPU上的加速效果确实突出但它的算子兼容性相对局限而且只适合英伟达的卡。Model-Optimizer在一开始就定了一个原则:不能绑死在特定硬件上。ONNX Runtime可以在CPU、GPU、甚至部分ARM设备上跑覆盖面广得多。真到了要在专用芯片上部署的场景ONNX也更容易转成各类厂商的模型格式。2.3 YAML配置驱动的流水线设计一个容易被忽略但又极度影响体验的设计是配置驱动。Model-Optimizer的每个优化任务都由一个YAML文件描述包含模型路径、优化步骤、数据集路径、剪枝比例、量化类型、蒸馏teacher模型地址等。这样做的理由是:团队成员不用读懂代码只要改配置就能跑通整个流程。举个例子一个典型的优化配置长这样:model: path: models/resnet50_imagenet.pth type: torchvision_classification input_size: [1, 3, 224, 224] pipeline: - step: prune method: bn_scale prune_ratio: 0.3 finetune_epochs: 20 - step: quantize method: ptq calibration_samples: 1000 backend: qnnpack - step: export format: onnx opset_version: 13每个step内部都是独立的类继承同一个OptimizerStep基类。基类负责加载模型、保存模型、记录日志这些公共逻辑;子类只实现run方法。这样新增一种优化手段比如加一个动态量化模块只需要新写一个继承类然后在pipeline配置里注册一下不需要改动主框架的代码。3. 实操把Model-Optimizer跑起来3.1 环境准备与依赖安装在开始之前先把环境列出来。Model-Optimizer本身不挑操作系统但我测试最充分的环境是Ubuntu 20.04 Python 3.9 CUDA 11.7。依赖库的安装顺序是有讲究的。推荐这样装:pip install torch1.13.1 torchvision0.14.1 --index-url https://download.pytorch.org/whl/cu117 pip install onnx1.13.1 onnxruntime1.14.0 pip install scipy scikit-learn pyyaml tqdm pip install nni2.10 # 如果需要用到NNI做超参搜索torch和torchvision版本一定要匹配否则后面导出ONNX时会出现奇怪的算子错误。onnxruntime装的是cpu版本还是gpu版本也要想清楚如果你在本地测试用CPU就够了生产环境再单独装GPU版本。千万别图省事直接pip install onnxruntime-gpu和CUDA版本不匹配的话日志里会报DLL加载错误。3.2 剪枝模块的实现细节剪枝模块的核心代码不长但每一行都有坑。我最初版本是这样实现BN层gamma筛选的:import torch import torch.nn as nn def prune_bn_scale(model, prune_ratio): bn_weights [] for name, module in model.named_modules(): if isinstance(module, nn.BatchNorm2d): bn_weights.append(module.weight.data.abs().clone()) bn_weights torch.cat(bn_weights) threshold torch.quantile(bn_weights, prune_ratio) prune_mask {} for name, module in model.named_modules(): if isinstance(module, nn.BatchNorm2d): mask module.weight.data.abs() threshold prune_mask[name] mask return prune_mask粗看没什么问题但实际运行时你会遇到两个痛点。第一个痛点是全局阈值 vs 逐层阈值。我在第一版用了全局阈值直接对整个模型的BN gamma做分位数裁剪。后来发现不同层的gamma分布差异极大浅层gamma普遍偏大深层gamma普遍偏小全局阈值容易把深层通道全部剪光。后来改成每层独立计算阈值每个BN层裁掉比例相同的通道。根据我的实测同样的30%稀疏度逐层剪枝的精度恢复速度比全局剪枝快很多。第二个痛点是实际剪枝要重新搭网络结构。光生成了mask没用你得根据mask重构卷积层。比如某个Conv2d原本输出通道是64mask保留了40个通道那这个Conv2d的输出通道数就得改成40同时下一层的Conv2d输入通道也要跟着改成40。这个联动关系如果处理不好模型结构就会对不上。Model-Optimizer里专门写了一个ChannelPruner类提交剪枝mask之后它会从被剪层的输出端一路传播到下一个消费层自动重建模型结构。3.3 量化模块的配置与参数说明量化部分我先说结论不管PTQ还是QAT校准数据集都需要覆盖真实业务分布。如果你拿ImageNet的几千张图去校准一个业务场景的检测模型得到的最小最大值范围往往偏得离谱。PTQ在Model-Optimizer里默认用的是Per-Channel量化因为卷积层每个输出通道的权重分布差异很大逐通道量化比逐张量量化的精度保持好一截。配置项里的qnnpack后端是针对CPU的它走的是MobileNet那套算子优化实测在ARM和x86 CPU上表现都不错。针对一次INT8量化后的精度回退我补充了QAT路径。QAT的关键技巧是在训练前端插入伪量化节点FakeQuantize让前向过程模拟量化的舍入误差反向传播时使用直通估计器这样梯度还能正常流动。Model-Optimizer里封装了一个prepare_qat函数内部用PyTorch的torch.quantization.prepare_qat然后接着做微调训练。QAT之后的模型精度通常能恢复到FP32模型的98%以上代价就是训练时间大概需要原来微调时长的1.5到3倍。我建议的量化路径是先用PTQ快速评估精度损失如果掉点超过2%再切到QAT。不要一上来就QAT时间成本不是开玩笑的。3.4 蒸馏模块的训练流程蒸馏模块我用的是典型的知识蒸馏配置。teacher模型权重加载到内存中student模型则随机初始化或者从预训练权重开始。训练损失由两部分组成loss alpha * ce_loss(student_logits, hard_labels) \ (1 - alpha) * kl_loss(student_logits, teacher_logits, T)alpha是硬标签损失和软标签损失的权重比例一般取0.5到0.7。T是温度参数作用是让teacher的softmax分布更“软”突出类别之间的相对关系。温度调高会让概率分布更平滑但太高也会引入太多噪声分类任务上T通常取4到8。蒸馏的一个容易被忽略的细节是teacher模型一定要设为eval模式。如果忘了切模式teacher的BN层还会跑批统计量的更新推理时分布会漂移student学到的软标签质量会明显下滑。另外teacher和student的输出尺寸不一定一致蒸馏一般要求两者的输出维度相同。如果student比teacher输出维度小可以在student后面接一个线性适配器或者只蒸馏中间层特征。Model-Optimizer里蒸馏模块不是必须的步骤。如果原来的模型本身不大只是推理环境苛刻那直接剪枝量化就够了。蒸馏更适合“大模型效果好但太大小模型效果好但难训练”的场景。4. 实测效果与调参记录4.1 在BERT和ResNet上的对比数据纸上谈兵没意思我把Model-Optimizer分别在两个典型模型上跑了一轮:一个ResNet50图像分类一个BERT-base的文本分类。测试机器CPU是Intel Xeon Platinum 8269CYGPU是NVIDIA T4量化目标统一是INT8 CPU推理。先看ResNet50的结果模型版本体积(FP32)延迟(CPU, ms)Top-1精度原始ResNet50102MB158.276.1%剪枝30% 微调72MB96.475.8%剪枝30% PTQ INT818MB42.775.3%剪枝30% QAT INT818MB42.575.9%表格里可以看出剪枝30%让体积降到原来的70%精度只掉0.3%在此基础上再做INT8量化体积进一步缩到18MBCPU延迟从158ms砍到42ms加速比接近3.7倍。QAT比PTQ在精度上多保住了0.6个点几乎和原始模型持平。再看BERT-baseSequence Length128批大小1模型版本体积延迟(CPU, ms)F1原始BERT-base438MB328.592.1%蒸馏到TinyBERT(4层)176MB121.390.8%蒸馏PTQ INT844MB56.990.6%蒸馏QAT INT844MB57.291.5%BERT这类Transformer在CPU上的INT8量化收益很大主要得益于Embedding和注意力算子在ONNX Runtime里的优化比较成熟。蒸馏到4层后体积减半再加上INT8量化整体体积只有原来的1/10延迟降到1/6F1只掉了1.5个点。在业务接受线上分类任务F1不低于90的前提下这个压缩幅度完全够用。4.2 参数选择背后的逻辑这些参数不是我拍脑袋定的每一个都经过实验推敲。第一个是剪枝比例的选择。ResNet50上我试过20%、30%、50%三档。20%剪完精度几乎无感知但加速有限50%剪完精度掉到73%以下微调20个epoch也难以恢复到74%。30%正好卡在精度和速度的甜点上。这个比例的确定不能照抄论文不同任务复杂度差异很大建议在25%到40%区间内做一轮网格搜索。第二个是量化校准集大小。PTQ的校准样本数量我设为1000张。试过500张、1000张、2000张发现500张时精度波动明显1000张以上就趋于平稳再多也只是增加校准时间。如果你对模型精度极度敏感可以把校准集加大到2000张但收益很边际。第三个是蒸馏温度T和alpha。在BERT蒸馏任务上T4、alpha0.5效果最好。T2时软标签太“锐”学到的类间关系信息少T8时噪声变大student模型训练不稳定损失曲线震荡。alpha也是同理硬标签损失占比太高会导致student过于自信预测分布生硬占比太低则连基本类别都学不牢。5. 常见问题与排查技巧5.1 剪枝后模型精度崩了怎么办剪枝完了精度暴跌这事每个人都遇到过不用慌。按照下面这个顺序排查第一先检查剪枝粒度。如果你做的是非结构化剪枝模型推理时其实用的是稀疏矩阵运算CPU和GPU不一定能享受到加速。而结构化剪枝后模型结构都变了必须重训才能恢复精度。如果你跳过了微调阶段直接部署精度当然是灾难。第二调整微调的超参数。剪枝后的模型相当于从一个较优的局部极小值点被强行拉到了另一个结构上此时学习率不能太大。我常用初始学习率2e-4配合Warmup逐步升高然后再余弦退火到0。训练轮次方面20个epoch是最低配置50个epoch能让精度恢复得更好但收益递减。第三可以只剪一部分敏感层。有些层比如第一层卷积和最后一层分类层对精度特别敏感我在Model-Optimizer里加了一个skip_layers配置剪枝时会自动跳过这些层。代价是压缩率稍微降低但精度稳定性大幅提升。5.2 量化模型在某些算子上的异常量化后模型能跑但个别算子报错或者输出异常这也是常见情况。比较典型的是LayerNorm和Softmax这类算子在部分ONNX Runtime版本上对INT8支持不完整。遇到这种问题我的做法是量化时对这些层使用“混合精度”策略——敏感层保持FP32计算其他层转INT8。在Model-Optimizer的量化配置里有个op_types_to_quantize白名单:quantize: op_types_to_quantize: [Conv, MatMul, Relu, Add] op_types_to_exclude: [LayerNormalization, Softmax]也就是说只量化Conv、MatMul这些对INT8友好的算子LayerNorm、Softmax这些动辄涉及精度敏感的层保留FP32。这样做的付出是延迟提升一些但换来了稳定性和精度。还有个小坑是动态量化 vs 静态量化。Transformer模型里Embedding层对输入长度敏感静态量化时序列长度固定还好一旦线上输入长度变化动态量化会更稳健。Model-Optimizer默认在BERT上用动态量化做Embedding静态量化做其余部分组合效果最稳。5.3 ONNX导出老失败的解决思路导出ONNX是另一个高频翻车点。常见报错是“Unsupported operator”或者“Missing opset version”。通常有两种解法:一种是升级opset版本。新版本ONNX支持更多算子但也可能引入新的行为变化。我在生产环境中固定opset13兼容性和算子支持都最均衡。opset17虽然有更新的特性但部分推理框架还没完全跟上。另一种是重写不支持的自定义算子。有些PyTorch的骚操作比如直接在forward里写了Python的for循环或者用到了torch.einsum导出时可能会拆得很碎。Model-Optimizer的导出模块里有一个onnx_simplifier步骤可以合并冗余节点和消除多余Transpose。如果简化后还不行就只好改写模型代码把自定义逻辑用标准算子拼接。6. 我踩过的坑和最后的建议模型优化这件事表面看是调参实际上是系统工程。我最深刻的体会是:永远要在优化之前先建立好完整的评估体系。Model-Optimizer里有一个很不起眼的evaluate模块但它是我最早的代码没有它后面的剪枝量化都无从谈起。每次做完一个step都要在同一份验证集上跑一遍精度和延迟结果记录到CSV里这样你才能判断每一步优化是不是真的有效。没有量化指标的优化都是自嗨。另外一个建议是:别追求极限压缩。Model-Optimizer在给业务方交付时我会强制要求精度回退上限是1%。一边是模型能小一点是一点另一边是业务方根本不会关心你的压缩率有多高人家只看准确率有没有跌。控制在指标范围内省下来的机器成本才是你真正的业绩;越界一个点前面省多少电费都堵不回来。最后分享一个小技巧把Model-Optimizer做成一个CLI工具而不是一堆散落的脚本。命令行接口的好处在于它强迫你规范化输入输出也方便别人接入CI流水线。只要配置好YAML任何模型进来都能一键跑完整套优化流程这才是工具的价值所在。至于后续还能扩展什么方向我觉得MVP最小可验证产品可以加自动化的结构搜索和更细粒度的敏感层分析不过我一般先把当前这套用熟再说。工具这事儿够用、稳定、不坑人就已经很厉害了。