把训练好的模型丢进生产环境,最头疼的往往不是精度,而是“跑不动”。我印象很深的一次,业务方要求推理时延压到5毫秒以内,服务器预算却一分不加。当时手里的模型是个标准的卷积网络,FP32权重占了40多MB显存,单次推理在GPU上跑出接近10毫秒。这个数字放到训练阶段完全不是问题,但在线推理场景里就是灾难。
那段时间我试了很多方案:换小模型、改部署框架、做算子融合,效果有,但都不够彻底。后来把目光放到模型压缩上,才算真正把问题解决。压缩这件事,主流路子就两条:剪枝和量化。剪枝负责把网络里冗余的参数和结构删掉,量化负责把计算精度从32位浮点降到8位整数,两个配合下来,模型体积、显存占用、推理时延都能有肉眼可见的改善。
但这两件事自己做起来很费劲。剪枝要考虑用什么标准挑权重、怎么控制稀疏度;量化要考虑校准集的选取、量化位宽、量化敏感层怎么办。一套流程手动搓下来,工作量不比训练一个模型小。后来我接触到NNI,这套工具把剪枝和量化从算法到训练再到导出都管起来了,整个过程顺了很多。这篇文章把我用NNI做模型压缩到部署的完整过程、参数选择和踩过的坑都记录下来,给同样在部署侧挣扎的朋友做个参考。
1. 部署为什么难,压缩为什么是刚需
1.1 模型在实验室和生产环境的落差
很多模型在训练阶段什么指标都好,一上生产环境就原形毕露。原因很简单,训练和部署是两个世界。
训练阶段你有A100、V100甚至更大显存的卡,不在乎模型多大、跑多快,只要精度够高就行。部署阶段换到普通的推理卡甚至CPU上,显存只有几个GB,时延要求几十毫秒以内,此时模型“又大又慢”就是原罪。尤其在大规模在线推理场景,单位成本是实打实核算的,很多团队每月算力账单里很大一部分花在模型膨胀带来的额外资源消耗上。
标准的部署优化手段,比如算子融合、图优化、更换推理引擎,本质上都在做同一件事:让计算更快一点。但模型本身的参数量和计算量没有变化,瘦死的骆驼比马大。当工程层面的优化空间被压榨得差不多时,就得回到模型本身动手,把参数减下去、把计算精度降下来,这就是剪枝和量化的价值。
1.2 剪枝和量化分别解决什么问题
打个比方,一个神经网络就像一套组织架构。剪枝做的事情是裁掉冗余的岗位和部门,组织结构更精干,处理流程更顺畅。量化做的事情是换一种更“省纸”的汇报方式,原来每份报告需要一页纸,现在用半页纸就能表达同样意思。
剪枝的核心价值在于减少计算量和参数量。网络训练过程中,大量神经元和权重对最终预测结果的贡献极小,甚至趋近于零。把这些“不干活”的部分去掉,模型变得更小更轻,但表达能力基本不变。量化则更激进,它不改变结构,而是改变数字的表示方式。原来用32位浮点数表示权重,现在改用8位整数表示,内存占用立刻缩小四倍,计算时因为数据带宽压力降低,速度也明显提升。
这两个手段不仅可以单独用,组合起来效果更佳。先剪枝把模型做小,再量化把精度降下来,最后部署时的模型尺寸和推理速度往往是数量级的提升。
1.3 为什么选NNI做这件事
做剪枝和量化其实有很多工具可选,但NNI有它特有的优势。首先它把“压缩”这件事做成了标准化流水线:定义压缩策略、执行压缩、微调、导出模型,每一步都有对应接口,不像自己写脚本那样散乱。
其次,NNI内置了多种剪枝算法和量化方法,覆盖了从入门到进阶的需求。剪枝方面有基于权重幅值、基于范数、基于泰勒展开等不同策略;量化方面支持量化感知训练、训练后量化等多种方式。你可以根据模型特点选择具体方法,不用什么都从零写。
再者,NNI的压缩流程是结构化的。它会记录哪些权重被剪掉、掩码是什么,导出的模型也是压缩后的真实模型,而不是剪完还得自己处理模型结构的半成品。这一点在实际项目中太重要了,曾经自己写剪枝脚本,导出的时候处理模型结构搞了很久,换成NNI之后一步到位。
2. NNI剪枝实操:先让网络瘦下来
2.1 剪枝到底剪掉了什么,怎么决定谁该被剪
剪枝有两种维度:结构化剪枝和非结构化剪枝。结构化剪枝直接移除整个通道、卷积核或层,模型结构变化明显,部署时能真实减少计算量;非结构化剪枝是抹掉单个权重,让权重矩阵变得稀疏,但结构不变,需要专门的稀疏计算库才能提速。
NNI里这两种都有支持,但我实际项目里更常用结构化思路去设计剪枝策略,因为最终落地部署时计算效率的提升更直接。NNI的剪枝器最终会在权重上生成掩码,把不重要的位置置零。置零的依据是权重本身的重要性,NNI提供了多种评估方式,比如L1范数修剪、L2范数修剪、基于激活的修剪等。
决定谁该被剪的核心是一套“重要性打分”机制。以最简单的L1范数剪枝为例,它会计算每个权重张量的绝对值,认为绝对值小的权重对结果影响弱,优先剪掉。如果剪枝粒度是通道级别,它会计算每个通道的范数,范数小的通道被整个置零。
这里有一个经验性的起步建议:稀疏度不要一开始就设得很高。我在做第一个项目时尝试过直接设0.8稀疏度,结果精度掉得没法看。后来做多了发现,从0.3到0.5之间起步,先看模型承受能力,能保持精度的前提下再逐步提高稀疏度,是更稳妥的方式。任何剪枝都不是无损的,稀疏度越高,精度损失越大,关键是要找到那个“损失小收益大”的平衡点。
2.2 NNI里怎么配置一次剪枝
NNI剪枝的入口很清晰,配置一次剪枝只需要做三件事:定义模型、写配置列表、实例化剪枝器。配置列表是NNI的核心,它告诉剪枝器剪什么、剪多少以及怎么剪。
import torch from torchvision.models import resnet18 from nni.compression.pytorch import L1NormPruner model = resnet18(pretrained=True) # 配置剪枝策略:对所有的卷积层做50%稀疏度的L1范数修剪 config_list = [ { "sparsity": 0.5, "op_types": ["Conv2d"] } ] pruner = L1NormPruner(model, config_list) pruner.compress()看到这段代码,你可能会问:对“所有的卷积层”统一用0.5稀疏度,真的合适吗?
实际训练中不同层的冗余程度差别很大,浅层捕捉的往往是基础纹理和边缘特征,冗余度较低;深层学到的语义特征更抽象,冗余度往往更高。针对这个问题,NNI支持更细粒度的配置,可以针对不同层设置不同稀疏度。
config_list = [ { "sparsity": 0.3, "op_types": ["Conv2d"], "op_names": ["conv1", "layer1.0.conv1"] }, { "sparsity": 0.6, "op_types": ["Conv2d"], "op_names": ["layer4.1.conv2"] } ]用op_names指定算子名字,可以为每一层量身定制剪枝方案。实际操作中,我一般会在配置文件里区分浅层和深层,浅层低稀疏度、深层高稀疏度,这样精度表现比全局一刀切稳定不少。
配置好之后调用compress(),剪枝就生效了。此时模型的前向过程已经带上掩码,被剪掉的权重不再参与计算。
2.3 剪枝之后的微调:不能省略的一步
剪枝完直接拿去部署,是我见过最多人犯的错误,我也踩过。L1NormPruner把权重置零之后,模型forward时会跳过这些权重,但模型本身的输出偏差已经产生了。剪枝破坏了原来训练时的参数分布,如果不做微调,精度往往惨不忍睹。
正确的流程是剪枝之后做一段短周期的训练,让剩余权重重新适应被剪掉的那部分参数的缺失。这个周期不需要像从头训练那样长,一般几个epoch到十几个epoch就够,学习率也建议比原始训练时调低一些。
# 剪枝后微调 optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9) criterion = torch.nn.CrossEntropyLoss() def finetune(model, dataloader, epochs=10): model.train() for epoch in range(epochs): for inputs, labels in dataloader: optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() print(f"epoch {epoch+1}, loss: {loss.item():.4f}") finetune(model, train_loader, epochs=10)微调完成后,模型的“剪枝痕迹”会被磨平,精度恢复到接近原始水平。这时候才轮到模型导出。
mask_dict, pruned_model_state = pruner.export_model( model_path="pruned_resnet18.pth", mask_path="mask.pth" )export_model是NNI剪枝流程里很关键的一个方法。它会重新整理模型状态字典,真正移除被剪掉的参数,生成一个压缩后的实体文件。这个导出的模型是可以直接加载部署的。
提示:export_model导出的state_dict是经过压缩后的模型形状。加载之前先用torch.load看一眼结构,避免部署时因为参数shape对不上而加载失败。
2.4 剪枝效果与精度损失实测
我在一个分类任务上用ResNet-18做过一次完整剪枝测试,用的就是上面的流程。原始模型参数量约1120万,剪枝稀疏度0.4之后,导出模型大小从43MB压缩到约26MB,参数量降到约680万。这个规模下精度只掉了0.3个百分点,基本可以忽略。
把稀疏度提高到0.6,模型体积进一步压缩到约19MB,但精度掉了1.7个百分点。这说明稀疏度存在一个临界点,超过之后精度损失会加速。到底能剪到多少,取决于模型本身冗余程度和任务复杂度,没有统一答案。
这里有个思考技巧:如果你的模型是从某个子网剪出来的,再剪的空间会很小;如果模型是那种大而全的网络,往往有惊喜。同样一个任务,用了一倍大的网络做剪枝,最终部署体积甚至可能比原模型还小,精度还能保持。
3. NNI量化实操:用更少的位跑出接近的精度
3.1 量化的底层逻辑:为什么8位数字也能算好
量化简单说就是换一种更省事的数字格式。深度学习训练和推理大多用FP32,即32位浮点数表示。一个8位整型INT8比FP32少了四倍空间,从存储和带宽角度看优势极其明显。
但你可能有个疑问:8位能表示的数字范围和精度都比32位差很多,怎么保证模型还能正常工作?
关键在于,神经网络对数值精度的宽容度比想象中高得多。训练完成后,权重虽然有微小扰动,但如果量化策略合理,这种扰动不会对最终的分类或检测结果有实质性改变。量化做得好,精度损失通常控制在1%以内,收益却是实打实的加速。
量化涉及两个核心参数:缩放因子scale和零点zero_point。它们负责把FP32的真实值映射到INT8的整数范围内。scale表示每个整数间隔代表多大的浮点数区间,zero_point表示浮点0对应到整数的哪个点。这两个参数的确定就是量化过程的核心工作。
NNI的量化器会自动帮你处理这些参数,但你是否理解它们的含义,决定了我后面能调出什么效果。校准和训练阶段选不好,量化误差就会累积起来。
3.2 两种量化路线:PTQ和QAT
模型量化有两条路线:训练后量化和量化感知训练。
训练后量化(PTQ)最直接:模型训练完了,拿一批样本做校准,统计权重和激活值的分布范围,然后确定scale和zero_point,把FP32权重变成INT8。校准只走几次前向过程,非常快。但因为对量化误差缺乏感知,精度损失可能比QAT大,特别当模型比较敏感时。
量化感知训练(QAT)则是在训练阶段就“模拟”量化的效果。前向传播时权重会被强制量化再反量化,相当于模型在训练过程中就适应了量化的“粗糙”,等真正导出INT8模型时,精度损失要比直接PTQ小很多。代价是要经历一个带量化模拟的训练周期,时间成本更高。
NNI对两种方式都有支持。实际项目里我会遵循一个选择逻辑:如果模型推理速度很吃紧但精度余量足,用PTQ快速试一轮;如果PTQ精度掉太多,再用QAT拉回来。一般我不会一上来就QAT,时间和算力都不划算。
3.3 QAT量化实操流程
以NNI的QATQuantizer为例,实现量化感知训练的核心代码大致如下:
import torch from nni.compression.quantization import QATQuantizer from nni.compression.pytorch.quantization.settings import QuantType config_list = [{ "quant_types": ["weight", "input"], "quant_bits": {"weight": 8, "input": 8}, "op_types": ["Conv2d", "Linear"] }] # 先加载一个已经微调好的剪枝模型 model.load_state_dict(torch.load("pruned_resnet18.pth")) optimizer = torch.optim.SGD(model.parameters(), lr=0.0005, momentum=0.9) qat_quantizer = QATQuantizer(model, config_list, optimizer, dummy_input) qat_quantizer.compress() # 量化感知训练 def qat_train(model, dataloader, epochs=5): model.train() for epoch in range(epochs): for inputs, labels in dataloader: optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() qat_train(model, train_loader, epochs=5) # 导出量化模型 qat_quantizer.export_model( model_path="quantized_pruned_resnet18.pth", calibration_info_path="calibration_info.pth" )QAT导入的dummy_input用于生成forward的输入张量形状,代码要跑得通,这一步不能省略。学习率我一般在原始训练学习率基础上再除以2或者除以10,因为量化感知训练本质上是个微调过程,步子迈大了容易把模型参数带偏。
导出时会同时生成calibration_info,里面包含量化参数。后续部署到推理引擎时需要加载这个文件,否则你还是拿到的是一堆FP32参数,量化白做了。
注意:导出模型的浮点权重和量化参数是分开保存的。在部署框架里,要么在转换时传入calibration_info,要么在推理脚本里手动加载,这一步很容易漏。
3.4 量化精度拉回技巧
量化后如果精度掉了,先别急着质疑工具。多数情况下是校准集选得不好或者量化敏感层没有单独处理。
校准集要贴近真实数据分布,数量不用多,几百张就够,但类别和场景要覆盖全。用训练集做校准其实是比较省事的做法,但如果训练集和线上数据分布有偏差,量化效果就会打折扣。我在实际中会把训练集里随机抽样一批和线上数据混合起来做校准,效果比单独用训练集更稳。
还有一种常见情况是模型中有个别层对量化特别敏感,一量化就崩。此时可以对这些层单独配置更高的bit位宽,比如权重用8位,敏感层用16位,其他层保持8位。NNI的config_list支持按op_names单独指定quant_bits,实现混合精度量化。
config_list = [ { "quant_types": ["weight", "input"], "quant_bits": {"weight": 8, "input": 8}, "op_types": ["Conv2d"], "op_names": ["layer1.0.conv1", "layer1.0.conv2"] }, { "quant_types": ["weight", "input"], "quant_bits": {"weight": 16, "input": 8}, "op_types": ["Conv2d"] } ]这段配置表达的意思是:前几层卷积权重保持16位精度,其他卷积层用8位。实际跑下来会发现,个别敏感层用高精度兜底,整体精度立刻回来大半,而体积只是略微增加。
4. 剪枝加量化组合部署:流程怎么搭最稳
4.1 先后顺序有讲究吗
剪枝和量化怎么组合,这个顺序在实操中很关键。
我的经验是:先剪枝,再量化。
原因主要有两个。第一,剪枝让模型变得更紧凑、参数分布更稀疏,这对后续量化是有利的。第二,先量化再剪枝会遇到一个麻烦,量化的权重分布已经变成离散的整数域,再做剪枝评估重要性时,那些基于浮点计算的评价指标会失效,剪枝效果会变差。
但有个细节需要注意:剪枝导出的模型在量化时最好重新微调,不要直接拿剪枝后的权重做PTQ。剪枝已经让模型精度有一些波动,再加上量化误差,两层叠加容易让精度滑出临界值。中间插入一小段量化感知训练,能明显消解这种叠加误差。
我的标准流程是:原始模型 -> 剪枝 -> 微调 -> QAT -> 导出量化模型 -> 部署。整个流程用NNI的接口串下来,代码逻辑是顺滑的。
4.2 模型体积、速度与精度的三方对比实测
用NNI走一遍完整流程后,我在一组测试数据上记录了对比如下:
| 模型状态 | 参数量 | 模型体积 | 推理时延(单次) | 精度 |
|---|---|---|---|---|
| 原始FP32 | 1120万 | 43MB | 9.8ms | 100%基准 |
| 剪枝后FP32(稀疏度0.4) | 680万 | 26MB | 7.5ms | 99.7%基准确 |
| 剪枝+量化INT8 | 680万 | 11MB | 4.2ms | 99.2%基准 |
| 剪枝+量化INT8(QAT微调后) | 680万 | 11MB | 4.2ms | 99.8%基准 |
数据说明两件事。第一,剪枝主要压缩体积和部分运算量,量化则把体积又压掉一半多,推理速度因为带宽降低和整数运算吞吐提升,有接近一倍的提升。第二,QAT微调这一步,确实把精度重新拉回来了,而且只用了5个epoch。
这组数据是单卡的实测结果。换成CPU推理,因为内存带宽对INT8的优势更敏感,量化带来的速度提升会更加夸张。有次在纯CPU环境测试,INT8模型比FP32模型快了两倍多。
4.3 导出后的模型怎么接进推理框架
NNI导出的是PyTorch格式的模型和量化参数,前端部署一般不会直接拿PyTorch跑,而是转换到推理引擎。这时有不少坑。
推理框架加载量化模型时需要知道每个算子的scale和zero_point,我一般直接将模型转为推理框架支持的格式,同时把calibration_info文件里的量化参数一并带入。大多数主流推理框架的量化转换工具都支持传入量化参数,不支持的框架,就只能手动在运行时加载参数了。
建议是:在导出的下一步立刻做格式转换和对比测试,早发现模型结构和算子兼容性问题,而不是等到部署环境里再排查。算子兼容性是量化部署最麻烦的问题,个别算子比如一些自定义激活函数在INT8下没有对应实现,这时需要在导入模型时代码里手动替换成量化框架支持的近似算子。
5. 踩坑记录与排查心得
5.1 常见问题速查表
| 现象 | 原因 | 解决方案 |
|---|---|---|
| 剪枝后精度暴跌 | 稀疏度太高,或剪后未微调 | 降低稀疏度,增加微调epoch数 |
| 量化后精度下降明显 | 校准集与真实分布偏差大 | 贴近线上数据选校准集,或改用QAT |
| 模型体积没变小 | 只compress未export_model | 调用pruner.export_model导出 |
| 加载模型时shape不匹配 | 忘记加载mask_dict影响模型结构 | 查看state_dict的shape,按需适配 |
| 推理时精度比预期低 | 部署引擎没有正确加载量化参数 | 检查calibration_info是否传入引擎 |
这些问题是项目中最常见的。很多是流程性错误,不是技术难点,提前看一遍完整流程就能规避大半。
5.2 一些实际的建议
做剪枝和量化,最需要用心的地方其实在前期。先搞清楚你的瓶颈在体积、时延还是功耗,再决定压缩策略的优先级。只追求小而不考虑精度,或者只追求快而不考虑质量,项目都跑不远。
我会建议你从一个小模型、一个简单任务开始,把NNI剪枝和量化的流程先完整跑通,记录每个阶段的模型大小、速度和精度变化,建立“压缩获得感”。有手感之后再上复杂模型,遇到问题排查起来也有参照系。
关于NNI版本带来的API差异,也值得多说两句。NNI版本迭代比较快,压缩模块的API在不同版本间有重新组织和命名的变动,比如有些剪枝器从旧算法的目录挪到了新目录,部分导入路径变化了。我写这篇文章使用的版本里,导入路径是from nni.compression.pytorch import这样形式的。如果你用的版本和这个不一致,导入报错时优先去官方文档确认当前版本的API签名,不要硬套旧版本的代码。
最后一个我反复踩过的坑:剪枝和量化的效果评估一定要在部署环境下做,不要在训练环境里看数字。训练环境里的张量运算和部署后的算子实现可能不一致,浮点计算在INT8算子上的细节差异也会影响最终精度。所以流程上尽早把导出模型放到目标推理环境里跑一轮验证,才能拿到真实收益数据。
我自己的感觉是,NNI把剪枝和量化的流程大大简化了,但模型压缩这事的核心思想——理解你的模型哪里冗余、哪里敏感,仍然要靠自己的经验去积累。动手做一版,记录数据,再调整,这个循环比任何理论都管用。