1. 为什么部署前必须折腾剪枝和量化真正开始碰模型部署之后我才发现训练阶段那句“模型越大、精度越高”的思路在线上基本走不通。带宽有限、显存有限、功耗有限一个动辄上百MB的模型塞进端侧光加载就要卡好几秒。于是“模型压缩”这一道工序几乎成了部署链路里的标配。而我做压缩时最常用的工具就是NNI。很多人以为NNI只是搜超参数的工具其实它自带的剪枝和量化模块已经帮我解决了好几个部署项目里的硬骨头。这篇文章就把我实际用NNI做剪枝和量化的过程完整拆一遍。1.1 部署场景里的硬约束模型的部署环境不是GPU机房随你挥霍的。端侧设备上限可能就是2GB内存服务端的显存预算也会被多个模型同时切分。训练时你关心的是Top-5有多高、收敛有多漂亮部署时你关心的是模型加载时间、单帧延迟、每秒处理多少请求、带宽够不够传一份更新包。一个300MB的图像模型在移动网络上分发给用户要下载半分钟这个体验基本是劝退级别的。还有功耗。同一个视觉模型在手机上跑一帧如果持续跑高帧率发热和掉电都很快。在部署链路里做剪枝和量化本质上是在“精度换资源”的杠杠上找平衡点。剪枝把冗余参数去掉模型体积和计算量同时下降量化把浮点运算换成定点运算既能减小体积又能利用推理引擎的INT8加速能力。两者不冲突而且通常配合使用效果最好。1.2 剪枝和量化到底动了模型的什么剪枝改变的是网络结构。它把那些对输出贡献很小、甚至趋近于零的权重或整个通道删掉相当于从原来的大部队里裁掉冗余人员。NNI里最常见的结构化剪枝是按通道删删完之后的模型是一个更窄的普通网络部署时不需要特殊支持。非结构化剪枝则是把散布在各处的单个权重置零精度常常更友好但稀疏矩阵要拿到实际加速必须部署引擎本身支持稀疏算子否则代码跑起来反而更慢。量化改变的是数值表示。它把FP32浮点权重和激活值换算成INT8整数看到的关键在于scale缩放系数和zero_point零点偏移。量化后的模型不再依赖FP32计算在支持INT8的CPU、GPU和NPU上往往是成倍的提速。更直白的类比量化像把RAW格式照片转成JPEG人眼看区别不大但存储和传输成本骤降。剪枝和量化分开做各有收益但项目里我更推荐“先剪后量化”的组合。先剪掉冗余结构再让量化过程小心翼翼地把数值压缩到8位最后得到的模型又小又快。1.3 为什么选NNI而不是自己手写模型压缩的工具并不少可以完全手写剪枝逻辑也可以用训练框架自带的接口。但手写的问题在于“算法库维护成本高”。一个剪枝算法从理论到可复现至少要搞定mask生成、层匹配、蒸馏/微调流程和后处理导出中间一步错了都可能把模型搞崩。NNI的价值是把常见剪枝器、量化器、训练微调工具统一封装API设计比很多个人实现规范得多而且支持PyTorch生态。对做部署的人来说把它当作一套“压缩工具箱”使用比从零手写要省力不少。关于NNI全称是Neural Network Intelligence一个开源的自动机器学习工具包。很多人不知道它自带一整套模型压缩能力从剪枝、蒸馏到量化都有现成封装。基于公开资料和实际使用情况NNI在2.x版本后的压缩模块逐步稳定API路径也清晰了不少。这里给个小提醒这套工具版本迭代较快如果跑着没问题就别随手升级固定一个版本长期用于部署能把很多不可控因素拦在门外。2. NNI剪枝的核心细节与算法选型NNI的剪枝流程先不看具体算法“mask”这个核心概念必须理解。Pruner不会直接删参数而是计算一张“保留/丢弃”的掩码表。比如一个卷积层有64个输出通道剪掉一半mask就是每个通道对应的0或1。真正的删除动作通常需要依靠ModelSpeedup把带mask的模型重新构图生成一个物理上更窄的网络。2.1 从API设计理解剪枝流程代码层面第一步是构造config_list。这是NNI统一的配置入口配置里写明“要给哪些层、剪多狠”。我的经验是指定sparsity稀疏度是好事但千万别一上来就对所有卷积层统一设0.8。不同层对剪枝的敏感度差异很大靠近输入的层承担底层特征提取一刀剪狠了后面全崩。from nni.compression.pytorch.pruning import L1NormPruner from nni.compression.pytorch.speedup import ModelSpeedup config_list [{ sparsity: 0.4, op_types: [Conv2d], exclude_op_names: [conv1] # 保留第一个卷积层 }] pruner L1NormPruner(model, config_list) _, masks pruner.compress() # 第一步拿到mask # 第二步真正从结构上删除被剪掉的通道 m ModelSpeedup(model, dummy_input, masks).speedup_model()这里有一个很容易踩的坑很多人以为调用compress之后就万事大吉直接拿原模型去推理。实际上mask只是标记了哪些通道要删模型的计算图还没有变化。如果只做稀疏推理可能影响不大但要部署到线上必须用ModelSpeedup做一次结构重排把被剪的通道物理删除。为什么保留第一个卷积层因为在很多视觉模型里输入层通道少、参数也少但它直接接收原始像素剪它往往精度损失大、收益小。一个通用策略是“先排除输入层和最后的全连接层优先从中间大通道数卷积层开刀”。2.2 常用剪枝算法怎么选NNI内置的剪枝器不止一种选错了算法效果天差地别。剪枝器原理适合场景注意事项L1NormPruner按权重L1范数排序范数小的先剪通用快速验证通道剪枝对分布集中的权重可能不够精细FPGMPruner以几何中位数为参考剪离中心最近的通道更深网络、规模较大的任务计算量稍大但对重要通道破坏小SlimPruner在训练中对BN的scale施加稀疏约束剪scale小的通道从零训练时就打算剪枝边训边剪需要修改训练流程不适合后处理直接剪MovementPruner按权重“移动量”排序关注训练中的变化趋势需要极致稀疏配合长时间重训训练成本高多为非结构化剪枝L1NormPruner是所有方案里最稳的起步选项。它的思路很直白权重绝对值小的通道对输出影响通常也小删掉它们风险最低。FPGM则更进一步考虑“即便权重范数小也不一定代表这条通道不重要”用几何中位数找真正冗余的簇中心深模型上精度保留更好。如果你的项目时间紧先用L1跑一版拿到baseline再做迭代。SlimPruner适合训练阶段就介入因为BN的scale本身就携带了通道重要性信息边训边剪能把稀疏约束织进网络参数里但它的流程改造比后处理剪枝麻烦适合有完整训练pipeline的团队。2.3 剪枝率不是拍脑袋敏感度分析与分层配置很多人上来就问“剪多少合适”这问题没有标准答案因为得看你模型冗余度、任务难度和数据规模。但有一个可行方法先做敏感度扫描。把某个层单独剪10%、20%、30%记录精度损失画出一条曲线。那些剪了20%精度纹丝不动的层就是理想的“肉层”剪5%精度就跳水的层必须降低稀疏度或者保留原样。具体配置上NNI支持按op_name细粒度配置。你完全可以对模型里不同层用不同sparsity。实际操作中我先用模型推理一遍收集每层输出分布再结合层参数量估算“可剪容量”最后写配置。比如对ResNet这类有多个stage的网络通常是后面stage的通道冗余高前面stage降精度更明显。我的一个经验值是整体目标50%稀疏度时把前面stage设0.2中间0.4最后stage 0.6这样整体稀疏度算下来仍然接近50%但精度损失往往比全局一刀切小一个量级。微调是必须的。剪完的模型不重训直接上线的做法除非剪得非常非常轻否则都是给自己挖坑。重新训练的学习率要比原始训练低大概在1e-4以下训练轮数至少是原始训练的20%到30%。轻量微调的意义在于让剩余参数“接管”被删通道的功能这个阶段如果发现精度回升缓慢可以适当调低剪枝率再重来而不是硬着头皮加epochs。3. NNI量化实操从浮点到定点做完剪枝下一步就是量化。NNI的量化模块里最有价值的是感知量化训练QAT而不是基础的训练后量化PTQ。PTQ不需要重新训练但模型越大、数据分布越复杂纯后校准的精度越难保证。QAT在训练过程中模拟量化误差让网络学会适应低精度表示精度保留通常更好。对于“已经压了一个大模型、还想再小一点”的项目QAT几乎是必选项。3.1 量化公式与核心参数先花两分钟把量化公式讲清楚后面配置才有底。非对称量化的核心公式是q round(x / scale) zero_point其中scale (max - min) / 255zero_point round(-min / scale)。这样浮点范围就被映射到0到255的整数区间。对称量化更简单不需要zero_point直接用[-127, 127]映射适合权重这种均值接近0的分布激活值分布往往偏移明显非对称更合适。NNI的config_list里你可以指定quant_types为weight还是input也可以只量化权重、不量化激活。我的经验是如果部署平台的INT8算子齐全权重量化激活量化一起开收益最大如果某些算子只支持权重量化那退一步也比完全不做强。量化bit选8是主流8位以下虽然体积更小但精度波动大一般项目不值得冒险。3.2 QAT感知量化实操要点NNI里配置QAT非常简单整体代码风格和剪枝很像还是熟悉的config_list结构。有一点要注意的是调用compress之前要先确认模型处于训练模式因为QAT本质上是一个训练过程模型内部的伪量化节点需要在forward和backward中参与更新。from nni.compression.pytorch.quantization import QATQuantizer qconfig [{ quant_types: [weight, input], quant_bits: {weight: 8, input: 8}, op_types: [Conv2d, Linear], exclude_op_names: [fc] # 按需排除 }] quantizer QATQuantizer(model, qconfig) model_qat, calibration_config quantizer.compress() model_qat.train()这里有个细节容易被忽略QAT过程中BNBatchNorm参数和量化scale是互相影响的。训练时BN的统计量会不断变化量化缩放因子也可能跟着跳。所以QAT的训练epochs不需要太长但学习率要更保守。如果你的模型里BN很多微调时最好先让BN统计量稳定下来再继续训或者直接在早期冻结BN的running_mean和running_var只更新权重和量化参数。这样做的逻辑是量化scale需要在一个相对稳定的数值分布上做校准如果BN统计量每轮都在变量化缩放因子就很难收敛。QAT之后不要急着直接部署伪量化模型伪量化只是模拟最终要继续导出成推理引擎能吃的格式。NNI处理后的模型支持torch.onnx.export导出导出时别忘了把模型切到eval模式。3.3 校准数据怎么准备如果走PTQ还需要一个校准数据集用来估计每层的动态范围。很多人直接拿训练集里的前几百张图这不一定靠谱。校准集有两个原则覆盖面要广不能只挑分类正确的样本数量不用多256到512张足够但必须能代表真实部署时的输入分布。最好从真实业务分布里采样比如线上日志里抽的图而不是训练集洗牌抽几张。QAT同样需要数据但它不是通过前向收集统计值就完事而是要把数据喂进训练循环里更新参数。所以准备QAT数据时必须带标签且类别均衡性也要照顾到免得量化模型在低频类别上翻车。我自己吃过的亏是用一整套干净训练集做QAT效果各方面都不错但一上线上真实数据精度就掉后来排查下来是校准数据的图像噪声和分布跟线上差异太大。从那以后我会刻意在QAT数据里混合一部分线上采集的样本让模型提前适应真实输入。4. 剪枝量化完整流程示例理论说完了直接上一条完整可跑的流程。以PyTorch NNI 2.x为例。任务就用一个轻量级图像分类Demo模型结构换成常规模块重点展示“怎么把剪枝、量化和导出串起来”。4.1 依赖安装与环境确认安装依赖时我的建议是固定版本尤其NNI这类迭代快的工具包API变动频繁。我常用的是2.10.1这个版本如果跑了不同版本个别导入路径可能要调整。这里不追求最新只求稳定可复现。pip install nni2.10.1 torch torchvision onnx onnxruntime装完先确认两件事PyTorch版本要和NNI兼容ONNX Runtime最好单独测一下int8算子支持。很多部署问题不是出在压缩阶段而是出在推理引擎对导出模型的支持上。可以在测试环境里先导出一个最小QAT模型用ONNX Runtime跑一遍确认没有不支持的算子再继续。4.2 第一步剪枝与结构重排以结构剪枝为例。先用预训练权重初始化一个普通卷积网络算出剪枝前baseline精度然后进剪枝流程。代码里我用一个简化模型占位实际使用时把它替换成你自己的网络就行。import torch import torch.nn as nn from nni.compression.pytorch.pruning import L1NormPruner from nni.compression.pytorch.speedup import ModelSpeedup class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 64, 3, padding1) self.conv2 nn.Conv2d(64, 128, 3, padding1) self.conv3 nn.Conv2d(128, 256, 3, padding1) self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Linear(256, 100) def forward(self, x): x torch.relu(self.conv1(x)) x torch.relu(self.conv2(x)) x torch.relu(self.conv3(x)) x self.pool(x) x torch.flatten(x, 1) return self.fc(x) model SimpleCNN() model.load_state_dict(torch.load(pretrained.pth)) model.eval() dummy_input torch.randn(1, 3, 224, 224) base_acc evaluate(model) # 剪枝前精度 config_list [{ sparsity: 0.5, op_types: [Conv2d], }] pruner L1NormPruner(model, config_list) _, masks pruner.compress() # 结构重排 model ModelSpeedup(model, dummy_input, masks).speedup_model()剪枝过程里的关键心得是dummy_input尺寸必须跟真实输入完全一致ModelSpeedup会依赖它做shape推导。有些自定义层不支持自动推导提前把这种层替换成标准卷积或去掉能省很多排查时间。还有一点speedup完成后的模型结构和PyTorch原生模型一致你可以继续走常规训练流程这一点让整个pipeline非常顺畅。4.3 第二步感知量化训练结构剪枝完成的模型是普通稠密网络可以直接做QAT。QAT之前先小规模微调几个epoch恢复精度再把量化器和微调一起做。注意量化训练阶段不要为了冲指标把数据增强拉满过强的增强会导致量化scale漂移。我在项目里通常会先跑一轮短训练观察loss曲线稳定后再继续往下走。from nni.compression.pytorch.quantization import QATQuantizer qconfig [{ quant_types: [weight, input], quant_bits: {weight: 8, input: 8}, op_types: [Conv2d, Linear] }] qat QATQuantizer(model, qconfig) model_qat, _ qat.compress() model_qat.train() optimizer torch.optim.SGD(model_qat.parameters(), lr1e-4, momentum0.9) for epoch in range(5): for x, y in train_loader: optimizer.zero_grad() out model_qat(x) loss criterion(out, y) loss.backward() optimizer.step()一个常见疑问是“QAT训练要不要从头训几十个epoch”。我的建议是适配场景如果模型是从公开权重拿来的5到10个epoch通常够了如果是自己从零训练的最好安排完整训练流程的30%作为QAT阶段而且学习率用余弦衰减更稳。实际操作里可以观察量化训练过程中精度曲线的波动如果一直上下震荡不收敛大概率是学习率太高或者数据增强太激进降一档再做。4.4 第三步导出并实测部署收益QAT完成后的模型在同一份测试集上验证一次然后导出ONNX。opset版本建议不低于12如果推理引擎支持动态batch配置dynamic_axes能减少后续对接工作量。这一步是整个流程的验收关口所有压缩工作的收益都要用最终导出的模型来衡量。model_qat.eval() torch.onnx.export( model_qat, dummy_input, model_deploy.onnx, opset_version13, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}} )拿到ONNX后用ONNX Runtime的CPU/GPU执行会话分别跑一遍对比延迟和模型体积。下面这张表是我在某个图像分类任务上的实测记录具体数字只是为了给你一个量级参考。影响速度的因素不只有参数量还有算子融合程度、缓存命中率、引擎是否启用INT8路径所以最终一定要以本机测出的数字为准。阶段模型体积推理延迟CPU单线程Top-1精度原始模型89MB220ms70.2%50%结构化剪枝46MB165ms69.9%剪枝 INT8量化22MB38ms69.4%这个例子里剪枝加量化之后模型体积缩到原来的四分之一单帧延迟直接降到原来的六分之一精度损失不到1个点。如果任务本身的冗余度更大收益还会更明显。但必须强调这不是天经地义的数值有些小模型压完可能速度不升反降原因我放到第五节细说。5. 踩坑记录与问题排查压缩这条链路我踩过的坑远比教程里写得多。整理几个高频问题权当速查表。5.1 剪枝掉精度剪枝后精度崩了优先排查三件事。一是稀疏度是否过高二是微调是否充分三是分层配置是否把关键层也剪了。这三个原因覆盖了大部分翻车现场。处理方法是先调低20%稀疏度跑一轮确认精度恢复后再逐步加大或者用FPGMPruner替换L1NormPruner对深层模型往往有改善。微调时如果精度卡住不动试着把BN层跑一个前向更新统计量再用更低的学习率接着训练。还有一个隐蔽原因剪枝前模型可能没有完全收敛那么剪枝后的微调是在一个不稳定的地基上修补效果自然差。所以我的习惯是剪枝前先确认模型已经在训练集上完全收敛最好再留出几个epoch只调BN让统计量稳定下来。5.2 导出ONNX模型通不过校验剪枝后的模型在PyTorch里正常一export就报错通常是自定义层或动态控制流惹的祸。NNI的ModelSpeedup对标准卷积、BN、ReLU这类原生层支持良好但对自定义op的shape推断很吃力。遇到这种报错先看日志里指向哪个layer能替换就替换不能替换就在导出前用torch.jit.trace单独追踪那段子图把问题范围缩小。opset版本也很关键。某些自定义算子需要更高opset才能支持但过高又可能与推理引擎不兼容。所以固定opset13是比较折中的选择既覆盖常见算子又避免导入时老版本引擎不认。5.3 量化后速度不升反降量化后模型变小但推理速度没改善甚至倒退别急着怀疑量化无用。先查推理引擎是否真正走了INT8路径。ONNX Runtime CPU下需要指定Execution Mode和关闭某些float fallbackGPU则需要对应厂商的INT8引擎。另一个常被忽略的因素是小模型上算子启动开销占大头量化带来的计算红利在极小模型上不明显。部署收益要用“每秒推理帧数”做基准而不是只看参数总量。我在测试一个轻量模型时遇到过类似情况模型本身就2MB量化之后体积减半但CPU推理延迟反而多了5ms原因就是INT8算子的调度开销盖过了计算提速。碰到这种情况更适合的做法是保持FP32配合通道剪枝来提速。5.4 量化后单测精度正常线上精度不对这个坑很磨人。量化模型在离线集上精度漂亮一到线上就拉胯多半是校准数据和真实输入分布不一致。我建议在准备QAT或PTQ校准数据时除了训练集一定要混入线上真实采样的数据最好覆盖各种图像质量、光照条件和类别比例。另一个角度是前处理如果和训练时不一致——比如归一化参数变了、输入分辨率变了量化的scale也会失真。检查一遍数据管线往往比调模型更快定位问题。5.5 NNI版本API变化带来的麻烦最后说一个工具本身的问题NNI版本切换时压缩API的包路径和返回结构可能会变网上旧教程容易让人卡壳。我的做法是把依赖版本写死在requirements里并且用一个独立的压缩脚本文档记录当前版本对应的API用法。项目交付后除非确有必要不轻易升级NNI版本。这样能让踩坑经验沉淀下来而不是每次升级重新踩一遍。6. 最后想说的几句这套流程跑通后我最大的感受是模型压缩不是把NNI里的函数挨个调一遍而是一个需要“设计、验证、反馈”的工程过程。先剪后量化、分层配置而不是一刀切、QAT阶段留足微调、导出后必须用真实部署引擎验证这几条是大方向上的关键。如果你准备做端侧部署或高吞吐服务端推理我非常建议拿一个已经收敛的模型从L1NormPruner加QAT的组合开始试起。第一次跑通后后续换模型、换任务就只是一遍遍套模板。压缩带来的收益很直观——更小的包体、更低的带宽、更快的响应每一分都是实实在在的成本和体验改善。踩过几次坑之后我还养成了一个习惯每次压缩完先把“模型体积、单帧延迟、Top-1精度”三组数字存成基线表格。下一次调整稀疏度或量化配置直接对照表格找最优解比凭借记忆来回试要高效得多。这个习惯我建议你也试试。