ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch模型量化实战:动态量化、静态量化与QAT加速推理

PyTorch模型量化实战:动态量化、静态量化与QAT加速推理

被问得最多的问题:模型推理太慢、显存不够,怎么办?很多人第一反应是换硬件、剪枝、蒸馏,但常常忽略一个性价比极高的手段——量化。PyTorch Python API 里的 Quantization 这个主题,就是专门解决这类工程问题的一整套工具,它能把 FP32 模型压到 INT8,在保持精度的同时,获得内存和速度上的巨大改善。这篇文章不打算复述官方文档,而是想聊聊我在实际项目中反复用到的三种量化方式,以及那些文档里不会明说的注意点。

如果你手里有一份已经训练好的模型,想在 CPU 上跑得更快、或者塞进内存受限的容器里,这篇内容可以直接照着操作。我会从量化原理入手,再分别介绍动态量化、静态量化和量化感知训练,最后把部署阶段容易翻车的坑一并列出。建议先收藏,实操时再打开对照。

1. 量化到底做了什么,为什么值得折腾

1.1 用一个缩放公式看懂量化

量化本质上是用低精度整数去近似连续浮点数的过程。在 PyTorch 里,最常用的是 INT8 量化,也就是把 FP32 的数值映射到 [-128, 127] 这个整数区间。核心公式并不复杂:

q = clamp(round(r / scale) + zero_point, q_min, q_max)

反量化的时候再乘回去:

r_approx = (q - zero_point) * scale

其中scale是缩放因子,决定浮点数值与整数刻度之间的比例;zero_point是零点偏移,对应浮点数值 0 时在量化整数域里的位置。举个直观的例子:一个 FP32 张量的取值范围是 [-2.0, 2.0],要映射到 INT8 的 256 个刻度,那么scale = 4.0 / 256 ≈ 0.0156,每个量化单位代表大约 0.0156 的浮点数值。反量化回来的数值必然存在一点误差,这就是量化误差的来源。

需要区分两种量化模式。对称量化里zero_point固定为 0,实现简单、计算更快,但对数值分布不对称的情况会浪费部分整数表示范围;非对称量化则通过zero_point把整数区间对准浮点数据的实际范围,能更好利用表示空间。PyTorch 的实际做法通常是:卷积和全连接层的权重用对称量化,激活值用非对称量化,这样兼顾精度和性能。

还有一个容易混淆的点是per-tensor与per-channel。per-tensor是对整个张量用一个scale和zero_point,实现简单但对通道间数值差异大的模型效果差;per-channel则对每个输出通道单独统计参数,精度更好,在 INT8 卷积中已经是标配。绝大多数 PyTorch 量化 API 默认给权重用per-channel,给激活用per-tensor。

1.2 量化红利从哪来:速度、内存、功耗

量化带来的收益是实打实的。第一是内存和带宽:INT8 只占 1 字节,FP32 占 4 字节,权重体积直接缩到四分之一。对于显存不足或内存带宽受限的推理场景,这往往比浮点计算本身更关键。第二是计算吞吐:现代 CPU 的 SIMD 指令集和 GPU 的 Tensor Core 对 INT8 都有专门加速通路,单位时钟周期内能完成更多算术操作。第三是功耗:低精度计算所需的数据搬运和算术单元复杂度都更低,边缘设备和移动端尤其在乎这一点。

但注意,“量化一定有加速”是个错觉。只有当算子的数值计算真正走在量化内核上时,红利才会兑现。如果你的模型里有大量自定义算子、频繁的reshape或复杂的控制流,访问内存和调度的开销可能会抵消 INT8 的优势。所以做量化之前一定要先做性能剖析,确认当前瓶颈到底是计算还是访存。

1.3 精度的代价为什么可控

量化毕竟是近似,为什么模型精度通常只损失一点点,甚至有时还会提升?这要从神经网络本身的鲁棒性说起。深度学习模型的中间特征之间通常存在大量冗余,参数和激活对微小扰动并不特别敏感。量化引入的误差相当于在模型里注入了一种特殊的“噪声”,只要这个噪声幅度不超过特征本身的扰动范围,损失就有限。

误差主要来自两个部分:取整误差和截断误差。取整误差是把实数近似为整数刻度时产生的,范围最多半个scale;截断误差则是数据超出可表示区间后直接被切到边界值造成的。截断误差往往更危险,因为离群点会被强行拉到 127 或 -128,导致分布改变。理解这两类误差之后,后面排查“为什么量化后精度崩了”就有了方向:要么是scale选得太大,要么是激活范围没覆盖好。

提示:量化的精度损失不是“模型变笨”,而是“信息有损地压缩”。如果能保证压缩后再展开的误差足够小,模型行为就不会明显变化。这也是为什么校准数据、Observer 的选择在量化里如此重要。

2. PyTorch 量化API全景:一个容易被中文文档忽略的版本问题

2.1 torch.quantization 与 torch.ao.quantization 的关系

打开 PyTorch 官方量化文档时,最常见的困惑是:一会儿看到torch.quantization,一会儿看到torch.ao.quantization,到底用哪个?简单说,早期 API 全部放在torch.quantization下,后来官方把 quantization 归入torch.ao(Automatic Optimization)子包,推荐写法变成了torch.ao.quantization。

在 PyTorch 2.x 中,torch.quantization和torch.ao.quantization大部分情况是兼容的,很多老代码依然能跑。但新项目建议直接使用torch.ao.quantization,因为官方在持续优化这块的命名和结构,新功能基本都落在torch.ao里,老路径以后会不会完全废弃也不好说。踩过一次旧代码升级的坑之后,我现在的习惯是统一import torch.ao.quantization as quant,代码更清晰,后面找人 review 也方便。

API 层面主要分为三大类:动态量化入口、静态量化入口(prepare/convert)、量化感知训练入口(prepare_qat/convert)。它们不是三套互不相干的东西,而是建立在同一个 Observer、QConfig、FakeQuantize 体系上。理解了这套体系,三种量化方式对你来说就不是背 API,而是同一个流程的不同变体。

2.2 三个重难点API:QuantStub / DeQuantStub / FakeQuantize

静态量化和 QAT 里,模型内部必须显式标注“在哪里量化、在哪里反量化”,这就要用到torch.ao.quantization.QuantStub和DeQuantStub。QuantStub负责把 FP32 输入转成量化张量,DeQuantStub负责把最后的量化结果转回 FP32 输出。新手容易忽略的是,这两个模块必须真实地写进模型forward里,而不是仅仅定义在类里。你可以在模型中控制量化作用的范围,比如只量化主干部分,输入输出层保持浮点。

FakeQuantize名字很直白,它“假装量化”但实际不改变数值类型。在 QAT 训练阶段,前向传播会先用低精度取整的规则模拟一遍量化误差,让模型感知到精度损失的存在;反向传播时却仍然用浮点梯度更新权重,从而让网络逐步适应量化。这就是 QAT 能比普通后量化训练出更高精度模型的核心原因。

这三个组件的关系可以这样理解:QuantStub/DeQuantStub决定了模型的“量化边界”,FakeQuantize决定了训练阶段“模拟误差的方式”,而后面要说的 Observer 决定了“量化参数怎么算出来”。

2.3 observer与QConfig:量化参数从哪来

量化参数scale和zero_point不是凭空产生的,而是从数据统计中算出来的。负责这件事的就是 Observer。常见的几个 Observer 分别是:

  • MinMaxObserver:直接统计张量最小值和最大值,最简单,但对离群点敏感;
  • MovingAverageMinMaxObserver:用滑动平均统计最小最大,对分布变化的激活更平滑;
  • HistogramObserver:基于直方图选择合适截断点,可以过滤部分离群点,静态量化里用得较多。

Observer 的可配置性很强,以MinMaxObserver为例,你可以设置量化比特数、是否per-channel、是否对称等。而在真实项目里,量化精度不高时我第一件事往往是换 Observer 配置,比如从MinMaxObserver换成HistogramObserver,经常能救回几个点的精度。

QConfig则是把这些 Observer 打包成一个描述“模型如何量化”的配置对象。它一般写成:

qconfig = torch.ao.quantization.QConfig( activation=torch.ao.quantization.default_fake_quant, weight=torch.ao.quantization.default_weight_fake_quant )

或者直接用官方预设的get_default_qconfig('fbgemm')。这里的fbgemm是后端字符串,对应 x86 CPU 上的 FBGEMM 内核;ARM 设备常用qnnpack,使用 Intel 优化库时可考虑onednn。选错后端不会报错,但可能拿不到预期性能,部署到不同设备前一定记得确认。

提示:Observer 在训练和推理阶段的统计策略不一样。训练模式下 Observer 会持续更新统计区间;eval 模式下 Observer 会冻结,使用已经统计好的参数。所以进行校准或评估时,模型的train/eval状态不能搞混,否则量化区间会一直在漂移。

3. 实测动态量化:一键压低模型体积,适合大模型与LSTM

3.1 适用场景与原理

动态量化是最省事的量化方式。它的思想是:把已经训练好的模型的权重提前量化成 INT8,但激活值仍然保持 FP32;每次前向的时候,根据当前层的实际激活分布临时计算出一个量化区间,把 FP32 激活量化到 INT8 后执行矩阵运算,输出再转回浮点。由于不需要提前收集统计信息,也就没有校准这一步,非常适合那些无法提供大量校准数据的场景。

动态量化对 RNN、LSTM、GRU 这类权重比较大、激活维度相对小的模型效果特别明显。文本模型、语音模型里很多结构是 LSTM 或 Transformer,这些模型往往受内存带宽限制,权重动态量化后访存量直接减少,推理延迟能下降许多。部分超大模型如果不方便重训,动态量化就相当于一个“几行代码的版本”。

3.2 动态量化实操代码

下面这段代码是我在文本分类模型上验证过的流程,几乎可以直接套用:

import torch import torch.nn as nn from torch.ao.quantization import quantize_dynamic model = nn.Sequential( nn.Linear(768, 256), nn.ReLU(), nn.Linear(256, 64), nn.Linear(64, 10), ) model.eval() quantized_model = quantize_dynamic( model, qconfig_spec={nn.Linear, nn.LSTM, nn.GRU}, dtype=torch.qint8 )

qconfig_spec可以传一个集合,表示只量化这些类型的模块;也可以传None,使用 PyTorch 默认支持列表。dtype通常用torch.qint8,CPU 推理场景这是最稳妥的选择。如果你看到官方文档里还提到torch.float16动态量化,它主要适用于某些特定后端,普通 CPU 上不建议优先尝试。

量化完成后,可以通过遍历模块确认权重是否真的变成了量化张量:

for name, module in quantized_model.named_modules(): if isinstance(module, nn.Linear): print(name, module.weight().dtype)

注意这里的module.weight()是方法调用,而不是module.weight。量化线性层的 weight 是一个torch.Tensor的子类,dtype 会显示为torch.qint8。如果你看到torch.float32,说明该层没有被量化。

我在一个大约 1 亿参数的 LSTM 模型上测试过,动态量化之后模型体积从约 400MB 降到 100MB 左右,推理速度提升接近 3 倍,且精度几乎无变化。对于部署到容器或边缘设备的场景,这个收益相当可观。

3.3 动态量化有哪些层不能碰

动态量化不是所有算子都支持。首先,nn.Embedding一般不会被量化,尤其在使用 one-hot 或稀疏更新时,量化反而可能拖慢速度。其次,自定义的nn.Module只要内部没有可支持的层,就会保持浮点。如果你想强行量化某个自定义层,需要把该层拆解成nn.Linear、nn.Conv2d等可量化模块,或者自己实现量化算子。

还有一点容易被忽略:动态量化对“已经量化过”的模型再次执行会报错或者产生不可预期行为,所以量化前一定要保证模型是 FP32 状态。如果你在实验里反复横跳,建议每次从原始 checkpoint 重新加载模型,再执行量化。

提示:动态量化的权重是离线量化的,但激活值是运行时实时量化。对于激活占主导的模型(比如某些轻量 CNN),提速可能很有限,这时候应该评估静态量化。

4. 静态量化:使用校准数据,获得更深层优化

4.1 为什么需要校准数据

静态量化与动态量化最大的区别,在于激活值的处理方式。静态量化会事先统计激活值的分布,确定好每个层的scale和zero_point,推理时直接使用这些固定参数,不再动态计算区间。好处是运行时节省了量化区间统计的开销,并且可以提前把层的融合、内存布局都优化好。代价是需要一个“校准”步骤。

校准不是训练,而是一次低成本的模型推理。你把一批具有代表性的真实输入喂给模型,让 Observer 统计每一层激活的数值范围。PyTorch 中prepare后的模型仍然以 FP32 运行,但所有可量化层都插入了 Observer,数据经过时会把范围记录在 Observer 内部。校准完成后执行convert,这些统计到的参数就会被固化到量化模型里。

校准数据的选择直接决定量化精度。有人图省事用随机数或训练数据的极端样本做校准,结果量化后精度暴跌,回头还怀疑是量化本身有问题。实数经验是:校准数据要覆盖真实推理场景中的典型分布,数量不用太多,几百到几千张通常就够。我的经验是至少覆盖每个类别的主要样本,且不要用数据增强后的图像,因为增强会引入分布偏移。

4.2 融合操作:Conv+BN+ReLU 为什么要先合并

做静态量化之前,PyTorch 官方流程都要求先做算子融合。最常见的是把Conv2d+BatchNorm2d+ReLU融合成一个Conv2dReLU量化算子。原因有两点:第一,减少算子之间的往返访存,融合后整个序列在一个内核里完成;第二,融合后把 BatchNorm 的缩放参数折叠进卷积权重,激活范围更稳定,量化误差自然更小。在量化之后做融合不是不行,但效果往往不如“量化前融合再重新校准”。

PyTorch 里融合需要你显式指定模块路径。如果你的模型是类似self.conv1、self.bn1、self.relu1的属性式写法,融合很简单:

model = quant.fuse_modules(model, [['conv1', 'bn1', 'relu1'], ['conv2', 'relu2']])

如果模型是用nn.Sequential堆的,路径就要写成字符串索引,比如['features.0', 'features.1', 'features.2']。踩过几次坑后,我建议在线性层、卷积层比较多的模型里尽量用带命名属性的写法,融和和后续调试都更省事。

融合需要满足相邻原则:Conv2d、BatchNorm2d、ReLU必须在模型结构中连续出现,中间不能隔着 Pooling 或 Dropout。像 MaxPool 这样的层通常不参与融合,可以在卷积加 ReLU 融合完成后单独保留。

4.3 静态量化完整示例(以小型CNN为例)

下面是一个可以直接运行的静态量化流程。我用了一个简单的小型 CNN,结构上包含两个卷积块和一个全连接层,并显式加入了QuantStub/DeQuantStub。这也是官方推荐的标准写法。

import torch import torch.nn as nn import torch.ao.quantization as quant class QuantCNN(nn.Module): def __init__(self): super().__init__() self.quant = quant.QuantStub() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.bn1 = nn.BatchNorm2d(16) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(2) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(2) self.fc = nn.Linear(32 * 8 * 8, 10) self.dequant = quant.DeQuantStub() def forward(self, x): x = self.quant(x) x = self.pool1(self.relu1(self.bn1(self.conv1(x)))) x = self.pool2(self.relu2(self.conv2(x))) x = x.view(x.size(0), -1) x = self.fc(x) x = self.dequant(x) return x model = QuantCNN().eval() model.qconfig = quant.get_default_qconfig("fbgemm") fused_model = quant.fuse_modules(model, [["conv1", "bn1", "relu1"], ["conv2", "relu2"]]) prepared_model = quant.prepare(fused_model, inplace=False) with torch.no_grad(): for images, _ in calib_loader: prepared_model(images) converted_model = quant.convert(prepared_model, inplace=False)

校准循环里要注意两点。第一,模型必须处于eval状态,确保 Observer 的统计不会被 BN 的 batch 统计干扰。第二,遍历校准数据的循环要使用torch.no_grad(),因为校准过程不需要反向传播,打开梯度只是徒增内存和时间开销。

convert完成后的模型可以正常forward推理,也可以在 CPU 上做基准测试。我自己实测,VGG 风格的小型 CNN 静态量化之后,单张图片推理延迟一般能降低 20%~50%,权重体积降为原来四分之一;精度通常在 99% 的 baseline 之上,少数模型几乎无损。

4.4 校准数据的数量与 Observer 选择

校准数据量并没有标准答案,但可以从误差趋势来判断。如果校准数据太少,Observer 可能只看到局部区间,激活范围的覆盖不完整;如果校准数据太多,各种离群样本会把范围拉得过大,导致有效量化精度下降。我的经验是 100~500 个 batch,单个 batch 大小与训练时一致,基本能覆盖主要分布。

当默认配置精度不理想时,可以尝试替换 Observer。比如默认的MinMaxObserver会把离群点也纳入区间,导致普通数值的量化精度变差;改用HistogramObserver就能按直方图分布选择更合理的截断点。一个快速实验是:

qconfig = torch.ao.quantization.QConfig( activation=torch.ao.quantization.HistogramObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine), weight=torch.ao.quantization.default_weight_observer ) model.qconfig = qconfig

这个配置对多数 CNN 都有正向效果。但 Observer 只是第一步,别指望所有精度问题靠它解决。

提示:静态量化最终推理强烈建议在 CPU 上验证。PyTorch 原生量化算子对 CPU 支持最成熟,GPU 上的 INT8 支持通常依赖 TensorRT 等外部推理引擎,别拿纯 PyTorch GPU 推理来评价量化效果。

5. QAT:把量化误差丢回训练里去学

5.1 原理:伪量化让网络自己适应

如果 PTQ(后训练量化)精度不够,就要上量化感知训练。QAT 的思路说起来很简单:在训练阶段就通过FakeQuantize模拟量化产生的取整误差,让模型在优化过程中主动适应这些误差。

具体实现上,QAT 的前向计算会先把权重和激活“伪量化”到低精度再转回浮点,因此 Loss 能反映出量化误差的影响;反向传播时,取整算子本身不可导,PyTorch 采用 STE(直通估计器)近似,把梯度直接绕过取整操作,继续更新浮点格式的权重。训练结束后,权重已经基本处在“适合被量化”的稳定状态,再做一次convert就能得到真正的 INT8 模型。

如果把 PTQ 看成“模型训练好后强行压缩”,QAT 就是“让模型在知道会被压缩的情况下训练”,后者自然更容易保住精度。它最大的劣势是训练开销。你需要准备标注好的数据、跑多个 epoch、调学习率,工程周期会明显拉长。所以我的建议是:先试试动态量化和静态量化,只有当精度下降不可接受时再考虑 QAT。

5.2 QAT 完整流程代码

QAT 流程与静态量化很相似,区别在于插入了prepare_qat和训练循环。下面这段代码基于一个普通的 CNN 分类模型:

model = build_pretrained_model() model.load_state_dict(torch.load("baseline.ckpt")) model.qconfig = quant.get_default_qat_qconfig("fbgemm") model = quant.fuse_modules_qat(model, [["conv1", "bn1", "relu1"], ["conv2", "relu2"]]) model = quant.prepare_qat(model.train(), inplace=False) optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) criterion = nn.CrossEntropyLoss() for epoch in range(epochs): for images, labels in train_loader: optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() model.eval() model = quant.convert(model, inplace=False)

有几个细节需要注意。prepare_qat传入的是model.train(),因为它需要开启 Observer 和 BN 的训练统计;但在训练的最后阶段,我们通常需要“冻结 BN”,让 BN 的运行均值稳定下来,否则convert后 BN 参数可能还在漂移。常见做法是在最后一个 epoch 里逐步降低学习率,或者训练结束后再用一小批校准数据走一次eval模式,让 BN 统计收敛到最终状态。

QAT 的学习率不能太大。量化感知训练的目的是微调模型适应量化误差,而不是从头训练,学习率太大会把预训练权重破坏掉。我通常从 baseline 训练学习率的十分之一开始。

5.3 QAT 与 PTQ 的选择决策建议

对比维度动态量化/PTQQAT
标注数据需求不需要/仅需少量校准数据需要完整训练数据
训练成本低,几分钟完成高,数小时到数天
精度表现多数模型可接受对敏感模型更可靠
实现复杂度简单,适合快速部署需要训练流程和调参
适用场景快速验证、模型较大精度敏感、量化后精度崩坏

如果你的模型是一个对量化特别敏感的网络,比如某些检测头、分割模型或极端稀疏网络,QAT 几乎是必经之路。但别在刚做完 PTQ 精度不理想时立刻上 QAT,先把 PTQ 的校准数据、Observer、融合都调一遍,因为这些环节对 QAT 同样有效,而且成本更低。

提示:QAT 完成后需要再次eval并做一次极小的校准循环,不是必须但很推荐。它能让 Observer 最终确定激活的边界值,避免convert时使用一个尚未收敛的统计结果。

6. 容易翻车的细节与排查经验

6.1 CPU/GPU差异与内存格式的多线恩怨

量化模型能不能在 GPU 上跑?这是个高频问题。PyTorch 原生的量化算子目前主要针对 CPU 推理优化,convert后的模型直接放到 CUDA 上经常得不到预期的 INT8 速度,有的算子甚至会回到 FP32 执行或者直接报错。因此,如果你的目标平台是 GPU,建议从开始就规划用 TensorRT 或其它 GPU 推理引擎做 INT8,而不是为难 PyTorch 原生 API。

CPU 推理时还有一个隐藏问题:量化的宽度与线程数。静态量化模型对张量内存布局有要求,PyTorch 内部会调用不同的后端实现,比如 FBGEMM 和 QNNPACK。如果只设置了qconfig而没指定后端库,某些算子可能走到速度较慢的 fallback 路径。我建议量化项目上线前做一次“后端字符串”对比实验,针对同一份模型分别用fbgemm、qnnpack和onednn构建,再在目标 CPU 上测延迟,选最快的。

6.2 算子白名单与黑名单:不支持的op会在convert时报错

模型里并非所有层都能量化。PyTorch 原生支持量化的常见算子包括nn.Conv2d、nn.Linear、nn.LSTM、nn.GRU、nn.ReLU、nn.AvgPool2d等。如果在convert阶段发现某个自定义层根本不支持量化,PyTorch 会抛异常,或者直接把该层留在浮点状态,而这往往不是我们想要的。

排查方法是逐一检查named_modules(),找到没有量化参数的层,再决定是替换为可量化实现,还是对那个子模块单独设置qconfig = None跳过。跳过是处理敏感自定义层的常见手段,例如第一层输入往往对量化敏感:

model.features[0].qconfig = None

这个操作要在prepare或prepare_qat之前设置,转换后的模型里该层会保持 FP32。还有一些算子如nn.Embedding,官方动态量化时一般不建议量化,因为收益有限且容易出问题,这种情况保持浮点反而是正确决策。

6.3 精度下降的排查清单

如果量化后精度明显下降,我建议按顺序排查,而不是盲目调参。

第一,先看校准数据质量。校准数据是否覆盖目标分布,是否混入特殊样本,数量是否足够。校准数据和训练数据分布不一致是静态量化精度崩掉最常见的原因。

第二,检查有没有做融合。未融合的Conv2d+BatchNorm2d+ReLU在量化后通常会有额外误差,融合之后误差往往明显下降。

第三,检查per-channel有没有开启。部分模型用per-tensor权重量化会损失较多,这时给权重配置per_channel=True,精度会有明显提升。

第四,定位敏感层。比较量化前后各层的输出分布或预测结果,找出差异最大的层,单独跳过或单独使用更高精度配置。我常用的方法是构造一组固定输入,分别用 FP32 模型和量化模型逐层输出,计算余弦相似度,差异大的层就是重点优化对象。

第五,考虑换 Observer。MinMaxObserver被离群点拖垮的情况经常出现,换成HistogramObserver往往能改善。前提是你对量化原理有一定理解,否则只是随机试参数。

6.4 模型保存、导出与后续部署问题

量化模型保存也有坑。虽然可以直接torch.save(model.state_dict(), ...),但量化模型里的权重可能是量化张量,重新加载时必须有对应的量化配置。更稳妥的做法是把整个量化后的模型用torch.jit.trace或torch.jit.script保存成 TorchScript:

scripted = torch.jit.trace(converted_model, example_input) scripted.save("model_int8.pt")

导出 ONNX 时同样要留意:convert后的模型是量化模型,ONNX 导出对 INT8 的支持与 PyTorch 版本、ONNX 算子集版本有关。别以为导出 ONNX 就万事大吉,INT8 ONNX 模型在推理引擎里还需要引擎侧的 INT8 支持和校准配置。我的项目经验是先导出 FP32 的 ONNX,再在不同引擎里做量化和校准,这样最可控。

还要提醒一点:如果要在 Python 服务里加载量化模型,加载后务必先跑一次model.eval(),否则 Observer 可能会继续统计新的数据范围,导致结果漂移。这个细节让我在线上环境耗了一晚上才排查出来。

6.5 torch.compile 与量化的关系

PyTorch 2.x 里torch.compile很流行,但它和量化的关系经常让人困惑。目前torch.compile对量化模型的支持并不是开箱即用的,尤其是静态量化和 QAT 流程中涉及 Observer、FakeQuantize 这些动态结构时,编译图可能会失效或退化。

如果你需要同时使用量化与torch.compile,建议先量化、再compile,或者反过来做小规模对比实验。不要默认两者叠加一定有收益,实测中经常出现编译后反而更慢或精度变化的情况。对于部署要求高的场景,我倾向于不用torch.compile,直接把量化模型导出 TorchScript 或 ONNX,用专门的推理引擎做优化。

7. 最后给新手的路线建议

我在实际项目里的习惯是:先写一个模型分析和精度对比脚本,分别跑一遍动态量化和静态量化,用可量化的指标(模型体积、CPU 延迟、精度变化)做决策。大部分项目到这一步就结束了,因为动态量化一行代码能拿到四分之三的收益,静态量化再把另一半收益吃进来。只有当精度掉得实在离谱,我才会上 QAT。

量化不是银弹,但它一定是最具性价比的部署优化手段之一。先把动态量化做了,再谈静态量化,最后才是 QAT,这是我推荐给所有人的安全路线。千万别一上来就冲着最高级的方案去,简单方案解决不了问题的时候,复杂方案才会显得有必要。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进