NNI 模型量化完全指南：从 QAT/PTQ 量化器到 TensorRT 推理加速

NNI 模型量化完全指南从 QAT/PTQ 量化器到 TensorRT 推理加速【免费下载链接】nniAn open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.项目地址: https://gitcode.com/gh_mirrors/nn/nni量化Quantization是 NNI 模型压缩体系中的核心能力之一通过降低权重与激活的数值表示位宽在几乎不损失精度的前提下显著减少模型存储体积与推理时延。本文以 NNI 官方量化文档docs/source/compression/quantization.rst为主线系统讲解量化的技术原理、五大量化器的算法差异、config_list 配置规范、QAT 与 PTQ 两种实战路径以及基于 TensorRT 的量化模型加速流程并结合 nni/compression/quantization 目录下的真实源码给出实现级解析。读完本文你将掌握在 NNI 中完成选量化器—写配置—训练/校准—导出—加速全链路实操能力。NNI 将模型压缩组织为预训练模型 → 剪枝 → 量化 → 压缩后模型的标准流水线量化位于压缩链的关键环节来源docs/img/compression_pipeline.png。剪枝降低参数维度保持 32 位宽量化压缩参数位宽32 位 → 8 位二者是两类互补的压缩策略来源docs/img/prune_quant.jpg。量化是什么从 FP32 到低比特表示在深度神经网络中模型权重的主流数值格式是 32 位浮点数FP32。量化指的是通过减少表示权重或激活所需的比特数来压缩模型的技术它能够同时降低推理的计算量与耗时。大量研究工作已经证明权重与激活可以用 8 位整数INT8表示而不会带来显著的精度损失更低比特位宽如 4/2/1 比特更是当前活跃的研究方向——这正是 NNI 量化模块nni/compression/quantization所覆盖的技术范围。需要区分两个关键概念量化Quantization把高精度数值映射到低比特离散值的过程。例如 FP32 的 4×4 矩阵量化为 8 位表示后维度不变但每个数值的存储与计算开销降为原来的 1/4。量化器QuantizerNNI 对量化算法实现的封装。NNI 的量化器都尽可能忠实复现其对应论文中的原始算法设计。从 NNI 源码结构看量化器的统一基类为Quantizer定义于 nni/compression/base/compressor.py所有量化算法QAT、DoReFa、BNN、LSQ、PTQ都继承自该类并共享一套量化目标空间 包装器 Evaluator 评估器的运行机制。NNI 五大量化器一览NNI 将量化算法的主体实现为量化器。下表给出了当前仓库中实现的量化器及其对应的论文出处点击名称即可跳转到对应量化器的详细接口文档docs/source/compression/quantizer.rst量化器名称算法简介对应论文QATQuantizer量化感知训练在训练前向过程中模拟量化舍入效应实现纯整数算术推理Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only InferenceJacob et al., CVPR 2018DoReFaQuantizer以低比特位宽同时量化权重、激活与梯度DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth GradientsBNNQuantizer二值神经网络权重与激活被约束为 1 或 -1Binarized Neural Networks: Training Deep Neural Networks with Weights and Activations Constrained to 1 or -1LsqQuantizer学习步长量化通过梯度训练量化步长scaleLearned Step Size QuantizationEsser et al., 2019PtqQuantizer训练后量化通过观测器在校准阶段收集量化信息无单一论文属于工程实践型算法从实现上看这五类量化器分别对应 nni/compression/quantization 目录下的qat_quantizer.py、dorefa_quantizer.py、bnn_quantizer.py、lsq_quantizer.py含lsqplus_quantizer.py与ptq_quantizer.py此外还有__init__.py统一导出全部类。其中 DoReFa 与 BNN 属于极低比特/二值量化路线源码中对激活模块的类型做了专门枚举如ACTIVATION_LIST覆盖 ReLU、GELU、Sigmoid 等常见激活函数见 nni/compression/quantization/dorefa_quantizer.py以支持对激活输出的量化。量化配置规范config_list 详解所有 NNI 量化器的第一个统一参数都是config_list——一个由 dict 组成的列表每个 dict 指明哪些模块需要被量化以及如何量化。NNI 将其完整规范定义在 docs/source/compression/config_list.rst下面梳理量化场景下的关键配置键。通用选择键锁定被量化模块配置键含义示例op_names待压缩模块的全限定名称列表若模块不存在会被忽略[backbone.layers.0.ffn, fc1]op_names_re用 Python 标准库re正则匹配模块名[.*conv.*]op_types继承自torch.nn.Module的类名列表仅匹配指定类型的模块未设置则所有类型都可选[Conv2d, Linear]exclude_op_names/exclude_op_names_re从待压缩集合中排除的模块名/正则[final_fc]exclude_op_types排除指定类型的模块[Dropout]若同时未设置op_names与op_names_re则匹配所有满足op_types的模块若op_types也未设置则默认匹配全部模块。target_names指定量化目标target_names是压缩目标的名称列表NNI 从设计上支持两类目标模块输入/输出若输入或输出是单个张量直接使用_input_输入与_output_输出。对于多输入场景可用_input_{position_index}或_input_{arg_name}精确定位例如前向函数def forward(self, x, y, z)中_input_0或_input_x都指向x计数时忽略self多输出场景同理可用_output_{position_index}或_output_{dict_key}。参数/缓冲区直接使用属性名如weight、bias。量化的典型配置是同时量化卷积/全连接层的输入、权重与输出例如target_names: [_input_, weight, _output_]。量化专属配置键除上述通用键外量化还使用以下专属配置键定义于 nni/compression/base/target_space.py 对应的属性解析逻辑中quant_dtype默认int8量化后的数据类型支持int和uint加量化比特数组合如int8、uint8、int4、int2、int1。源码中按quant_dtype.startswith(int)/startswith(uint)解析出quant_bit并据此计算量化区间[qmin, qmax]见 nni/compression/base/target_space.py。注意int1与 BNN 的 ±1 二值化并非同一概念BNN 仅把数值约束到 1/-1quant_dtype对 BNNQuantizer 不生效源码中会给出警告。quant_schemeaffine或symmetric量化方案。未设置时由量化器自行选择多数量化器默认应用symmetric对称量化。二者的 scale / zero point 计算差异直接体现在源码的compute_scale_zpnni/compression/quantization/ptq_quantizer.py与update_scale_zpnni/compression/quantization/qat_quantizer.py中symmetric对称scale abs_max / ((qmax - qmin) / 2)zero point 固定为(qmax qmin 1) // 2PyTorch 中 qint8 的 zp 为 0quint8 的 zp 为 128。affine非对称scale (tracked_max - tracked_min) / (qmax - qmin)zero_point qmin - round(tracked_min / scale)并 clamp 到[qmin, qmax]。两种方案都会把 min/max 区间扩展到包含 0以保证 0 是可精确表示的数值。granularity量化粒度控制目标量化的粒度默认整张张量共享同一组 scale 与 zero point。支持两种取值per_channel每个输出通道拥有独立的 scale 与 zero point整数列表以列表为 block size每个 block 拥有独立的 scale 与 zero point块量化。注意BNNQuantizer不支持 granularity其_scaler必须为 None否则直接抛ValueError见 nni/compression/quantization/bnn_quantizer.py。fuse_names可选参数类型为List[Tuple[str, ...]]。每个 tuple 定义需要融合进第一个模块的模块列表tuple 中第一个模块名必须出现在op_names或op_names_re中。该键用于把 BatchNorm 等算子融合进卷积层再量化避免逐层模拟量化误差累积。config_list 的覆盖规则config_list 中每个子配置是一个 dict各键的作用域仅限于该子配置内部。若多个子配置同时命中同一层后出现的子配置会覆盖先前的配置。利用这一规则可以像官方 quickstart 那样用两条子配置分别描述卷积/全连接层量化 input/weight/output与激活层仅量化 output。实战路径一量化感知训练QAT量化感知训练Quantization-Aware Training把量化效果在训练的前向传播中模拟出来权重在参与卷积前先被量化若该层使用了 BatchNorm其参数会先折叠进权重再做量化激活在推理时会被量化的位置如卷积/全连接层输出经过激活函数之后、残差连接相加处进行量化。反向传播仍按常规进行所有权重与偏置依然以浮点形式存储从而可以被小幅微调以补偿量化误差。NNI 的QATQuantizer完整实现了这一算法其类文档直接引用了原论文见 nni/compression/quantization/qat_quantizer.py。QATQuantizer 的参数参数说明model待量化模型config_list量化配置列表指明哪些模块被量化、如何量化evaluator评估器封装训练循环/优化器/训练步函数quant_start_stepQAT 开始前的预热warmup步数默认 0从源码看quant_start_step的机制很直观在预热阶段所有量化目标空间的apply_method被设为bypass跳过量化模拟当训练步数current_step quant_start_step时切换为qat_clamp_roundnni/compression/quantization/qat_quantizer.pyscale 与 zero point 也只在current_step quant_start_step之后才随训练更新同文件 L115-L148。此外 QAT 通过track_min_max_val以 0.99 的衰减系数做指数移动平均EMA来跟踪张量的 min/max见同文件 L183-L208scale/zero point 即由 EMA 后的极值计算而来。完整 QAT 实战代码MNIST官方 quickstart 教程docs/source/tutorials/quantization_quick_start.rst源码见 examples/tutorials/quantization_quick_start.py以 MNIST 为例演示完整流程。先定义一个带卷积、池化、ReLU6 与全连接的小模型并准备好训练/评估 DataLoaderimport torch import torch.nn.functional as F from torch.optim import SGD from torch.utils.data import DataLoader from torchvision import transforms from torchvision.datasets import MNIST MNIST(rootdata/mnist, trainTrue, downloadTrue) MNIST(rootdata/mnist, trainFalse, downloadTrue) transform transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) train_dataloader DataLoader(MNIST(rootdata/mnist, trainTrue, transformtransform), batch_size64) test_dataloader DataLoader(MNIST(rootdata/mnist, trainFalse, transformtransform), batch_size1000) class Mnist(torch.nn.Module): def __init__(self): super().__init__() self.conv1 torch.nn.Conv2d(1, 20, 5, 1) self.conv2 torch.nn.Conv2d(20, 50, 5, 1) self.fc1 torch.nn.Linear(4 * 4 * 50, 500) self.fc2 torch.nn.Linear(500, 10) self.relu1 torch.nn.ReLU6() self.relu2 torch.nn.ReLU6() self.relu3 torch.nn.ReLU6() self.max_pool1 torch.nn.MaxPool2d(2, 2) self.max_pool2 torch.nn.MaxPool2d(2, 2) self.batchnorm1 torch.nn.BatchNorm2d(20) def forward(self, x): x self.relu1(self.batchnorm1(self.conv1(x))) x self.max_pool1(x) x self.relu2(self.conv2(x)) x self.max_pool2(x) x x.view(-1, 4 * 4 * 50) x self.relu3(self.fc1(x)) x self.fc2(x) return F.log_softmax(x, dim1)然后定义训练步、训练循环与评估函数并在 MNIST 上预训练 5 个 epochdevice cuda:0 if torch.cuda.is_available() else cpu def training_step(batch, model) - torch.Tensor: x, y batch[0].to(device), batch[1].to(device) return F.nll_loss(model(x), y) def training_model(model, optimizer, training_step, schedulerNone, max_stepsNone, max_epochsNone): model.train() max_epochs max_epochs if max_epochs else 1 if max_steps is None else 100 current_steps 0 for epoch in range(max_epochs): for batch in train_dataloader: optimizer.zero_grad() loss training_step(batch, model) loss.backward() optimizer.step() current_steps 1 if max_steps and current_steps max_steps: return if scheduler is not None: scheduler.step() def evaluating_model(model): model.eval() correct 0 with torch.no_grad(): for x, y in test_dataloader: x, y x.to(device), y.to(device) preds torch.argmax(model(x), dim1) correct preds.eq(y.view_as(preds)).sum().item() return correct / len(test_dataloader.dataset) model Mnist().to(device) optimizer SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) training_model(model, optimizer, training_step, None, None, 5) print(Acc.:, evaluating_model(model))接下来是核心的量化步骤。需要先把优化器用nni.trace包装并构造TorchEvaluator交给量化器NNI 借此在优化器 step 上注入量化逻辑import nni from nni.compression.quantization import QATQuantizer from nni.compression.utils import TorchEvaluator optimizer nni.trace(SGD)(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) evaluator TorchEvaluator(training_model, optimizer, training_step) config_list [{ op_names: [conv1, conv2, fc1, fc2], target_names: [_input_, weight, _output_], quant_dtype: int8, quant_scheme: affine, granularity: default, }, { op_names: [relu1, relu2, relu3], target_names: [_output_], quant_dtype: int8, quant_scheme: affine, granularity: default, }] quantizer QATQuantizer(model, config_list, evaluator, len(train_dataloader)) real_input next(iter(train_dataloader))[0].to(device) quantizer.track_forward(real_input) _, calibration_config quantizer.compress(None, max_epochs5) print(calibration_config) print(quantized Acc.:, evaluating_model(model))这里的要点config_list 用两条子配置分别覆盖卷积/全连接层量化输入、权重、输出与 ReLU6 激活层仅量化输出激活输出按 NNI 的约定使用_output_目标名QATQuantizer的第四个参数len(train_dataloader)即quant_start_step——让模型先预热训练一个 epoch 再开启量化模拟quantizer.track_forward(real_input)用真实输入执行一次前向以初始化 scale / zero point 等跟踪缓冲源码见 nni/compression/quantization/qat_quantizer.pyquantizer.compress(None, max_epochs5)返回(model, calibration_config)calibration_config是一个嵌套 dict记录了每个模块每个量化目标的scale、zero_point、quant_dtype、quant_scheme、quant_bits以及追踪到的tracked_max/tracked_min。官方教程的实测输出显示FP32 模型在 MNIST 上准确率约 0.9906而经过 5 个 epoch QAT 训练后量化模型准确率约 0.9912精度基本持平。教程对应的完整可运行脚本位于 examples/tutorials/quantization_quick_start.py另一个结构相近、可直接执行的示例见 examples/compression/quantization/qat_example.py。仓库中的其余 QAT 变体示例如 examples/compression/quantization/bnn_example.py、dorefa_example.py、lsq_example.py与本文 QAT 用法完全一致仅需替换量化器类名。实战路径二训练后量化PTQ与校准PtqQuantizer是 NNI 的训练后量化实现在校准阶段用观测器observer收集量化信息无需反向传播训练适合数据量有限或不想改动训练流程的场景。其核心流程见 nni/compression/quantization/ptq_quantizer.py先以bypass方式不量化用评估数据跑一遍模型track_min_max_val记录每个量化目标的 min/max同文件 L75-L96评估结束后调用update_scale_zp依据compute_scale_zp计算出 scale / zero point同文件 L98-L102、L141-L161将apply_method切换为clamp_round真正进入量化模拟可选is_bias_correctionTrue开启偏置校正通过比较量化前后模型输出统计量来修正 bias校正后会重新计算 scale / zero point同文件 L122-L138。一个典型的 PTQ 配置如下来自 docs/source/tutorials/quantization_speedup.rstfrom nni.compression.quantization import PtqQuantizer from nni.compression.utils import TorchEvaluator def eval_for_calibration(model): evaluate(model, data_loader, neval_batches1, devicedevice) # 校准数据 dummy_input torch.Tensor(64, 3, 224, 224).to(device) predict_func TorchEvaluator(predicting_funceval_for_calibration, dummy_inputdummy_input) config_list [{ quant_types: [input, weight, output], quant_bits: {input: 8, weight: 8, output: 8}, quant_dtype: int, quant_scheme: per_tensor_symmetric, op_types: [default], }] quantizer PtqQuantizer(model, config_list, predict_func, True) quantizer.compress() calibration_config quantizer.export_model() print(calibration_config)注意该配置片段沿用了旧版 API 的quant_types/quant_bits/quant_scheme: per_tensor_symmetric写法教程中skip_exec场景下未实际执行当前版本的config_list推荐使用target_namesquant_dtypequant_schemeaffine/symmetricgranularity的写法与本文 QAT 示例保持一致。PtqQuantizer通过TorchEvaluator(predicting_func..., dummy_input...)在校准阶段用评估数据做前向统计export_model()导出与 QAT 相同结构的calibration_config。从量化模型到真实加速TensorRT 推理引擎需要特别强调的是QAT / PTQ 等量化算法通常以模拟方式工作——用 FP32 算子模拟低比特计算的效果张量被量化到目标比特数后再反量化回 FP32因此这类量化模型本身并不会带来时延降低。要让量化模型真正跑在低比特算子加速上必须有一个 speedup 阶段。NNI 当前通过TensorRT 推理引擎实现量化模型的真实加速见 docs/source/tutorials/quantization_speedup.rst。NNI 的加速流程分两步1) 将带量化权重与配置的模型转换为 ONNX 格式2) 将 ONNX 模型喂给 TensorRT 生成推理引擎用于低时延推理。对应实现位于 nni/compression/quantization_speedup 目录frontend_to_onnx.py负责 PyTorch→ONNX 转换integrated_tensorrt.py封装 TensorRT 引擎构建。环境准备使用 TensorRT 加速量化模型时强烈建议使用 NVIDIA 提供的 PyTorch Docker 镜像如已测试过的nvcr.io/nvidia/pytorch:22.09-py3。启动容器后需安装nni3.0、pytorch_lightning与pycudanvidia-docker run -it nvcr.io/nvidia/pytorch:22.09-py3 pip install nni3.0 pytorch_lightning pycuda模式一直接利用 TensorRT 的训练后量化带校准数据在该模式下QAT / LSQ 等量化感知训练算法只负责把模型权重调整得更利于量化最后一步量化交给 TensorRT 自带的训练后量化完成。示例流程以 MobileNetV2 与 ImageNet 为例from nni.compression.quantization_speedup.calibrator import Calibrator from nni.compression.quantization_speedup import ModelSpeedupTensorRT # 1. 准备 128 个样本的校准数据TensorRT 按 batch_size64 处理 calib_data ... # shape 为 (128, 3, 224, 224) 的 numpy 数组 calib Calibrator(calib_data, data/calib_cache_file.cache, batch_size64) # 2. 加载 FP32 模型并直接用 calibrator 加速 model MobileNetV2() state_dict torch.load(mobilenet_pretrained_float.pth) model.load_state_dict(state_dict) model.eval() engine ModelSpeedupTensorRT(model, input_shape(64, 3, 224, 224)) engine.compress_with_calibrator(calib)随后可用engine.inference(image)获得加速后的推理输出与耗时time_span并统计 Top-1 / Top-5 准确率。模式二把 TensorRT 当作纯加速后端该模式不使用 TensorRT 内部的训练后量化而是把我们自己量化得到的位宽与张量取值范围喂给 TensorRT内部通过trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS配置实现。流程为先用PtqQuantizer完成量化并导出calibration_config再重新实例化原始模型因为 BN 折叠已改变模型结构与权重而 TensorRT 会自行做 BN 折叠所以需要喂给它原始模型最后按配置构建引擎from nni.compression.quantization import PtqQuantizer from nni.compression.quantization_speedup import ModelSpeedupTensorRT # 用 PtqQuantizer 量化并得到 calibration_config quantizer PtqQuantizer(model, config_list, predict_func, True) quantizer.compress() calibration_config quantizer.export_model() # 重新实例化原始模型按量化配置构建 TensorRT 引擎 model MobileNetV2() model.load_state_dict(torch.load(mobilenet_pretrained_float.pth)) model.eval() engine ModelSpeedupTensorRT(model, input_shape(64, 3, 224, 224), configcalibration_config) engine.compress()两种模式的选择原则若希望省事、直接利用 TensorRT 成熟的 PTQ 能力选模式一若已有 QAT/PTQ 产出的量化配置、希望完全遵循自己的量化方案选模式二。相关资源与延伸阅读量化概述本文主体文档 docs/source/compression/quantization.rst量化器索引表 docs/source/compression/quantizer.rst配置规范 docs/source/compression/config_list.rst快速上手教程 docs/source/tutorials/quantization_quick_start.rst可运行脚本 examples/tutorials/quantization_quick_start.py加速教程 docs/source/tutorials/quantization_speedup.rst量化器源码 nni/compression/quantizationqat_quantizer.py/ptq_quantizer.py/bnn_quantizer.py/dorefa_quantizer.py/lsq_quantizer.py加速后端源码 nni/compression/quantization_speedup可运行示例 examples/compression/quantizationqat_example.py/bnn_example.py/dorefa_example.py/lsq_example.py/ptq_example.py/qat_example.py结合仓库中的 docs/source/compression/toctree.rst 与 docs/source/compression/toctree_quantization.rst 可以看到NNI 还提供 BERT/GLUE 上的 Transformer 量化教程docs/source/tutorials/quantization_bert_glue.rst以及模型蒸馏、剪枝与量化融合压缩等能力量化与剪枝可组合进同一条压缩流水线见 docs/source/compression/overview.rst适合作为本文之后的进阶路线。【免费下载链接】nniAn open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.项目地址: https://gitcode.com/gh_mirrors/nn/nni创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考