企业官网建设流程全解析

PyTorch-Quantization 量化工具包实战指南：从模拟量化训练到 TensorRT INT8 部署

【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT

PyTorch-Quantization 是 NVIDIA TensorRT 开源仓库tools/pytorch-quantization下的一套 PyTorch 量化工具包（仓库内版本 2.2.1，见 VERSION），用于以“模拟量化（Simulated Quantization）”的方式训练与评估量化模型，其产物可导出为 ONNX 并被 TensorRT 8.0 及更高版本导入执行，充分发挥整数 Tensor Core 的高性能推理能力。读完本文，你将掌握该工具包的安装方式、QuantDescriptor/TensorQuantizer核心概念、自动与手动量化、四种校准方法、量化感知训练（QAT）以及 ONNX 导出部署到 TensorRT 的完整链路。

维护状态说明：据 README.md 开头声明，PyTorch-Quantization 的后续开发已移交 NVIDIA 的 TensorRT Model Optimizer，官方建议新项目迁移至 Model Optimizer 以获取最新的量化与压缩能力；本工具包代码仍保留可用，但不再继续演进。本文内容以当前仓库实际代码与文档为准。

1. 工具包定位与仓库结构

PyTorch-Quantization 的核心定位（摘自 README.md）可以概括为三点：

• 模拟量化训练与评估：量化既可以自动加入模型，也可以手动加入，以便针对精度与性能进行调优；
• 与整数 Tensor Core 兼容：量化模型可直接适配 NVIDIA 的高性能整数内核；
• ONNX 导出 + TensorRT 部署：量化后的模型可导出为 ONNX，由 TensorRT 8.0 及以后版本导入并执行。

从源码结构看，工具包分为以下几个层次（各模块均有对应测试，见 tests/ 目录）：

2. 安装与运行环境

README 提供了三种安装途径，覆盖从快速试用到源码定制的全部场景。

2.1 二进制安装（pip）

pip install pytorch-quantization --extra-index-url https://pypi.ngc.nvidia.com

该命令从 NGC 的 PyPI 镜像拉取预编译包，适合无需改动源码的常规使用。

2.2 源码安装

获取仓库源码并进入工具包目录：

git clone <TensorRT 仓库地址> cd tools/pytorch-quantization

先安装 PyTorch 与依赖：

pip install -r requirements.txt # for CUDA 10.2 users pip install torch>=1.9.1 # for CUDA 11.1 users pip install torch>=1.9.1+cu111

然后构建并安装：

# Python version >= 3.7, GCC version >= 5.4 required python setup.py install

几点环境细节（源自仓库实际文件）：

• requirements.txt 声明的运行依赖为：numpy、absl-py>=0.7.0、scipy、sphinx_glpi_theme、prettytable、pyyaml；其中scipy服务于直方图校准中的熵计算（见 calib/histogram.py 中对scipy.stats.entropy的引用）；
• setup.py 使用 PyTorch 的torch.utils.cpp_extension（CppExtension/CUDAExtension/BuildExtension）编译src/下的 C++ 与 CUDA 扩展，因此要求本机具备可用的 CUDA 编译工具链；
• setup.py会在构建时依据 VERSION 生成pytorch_quantization/version.py，当前仓库版本为2.2.1。

2.3 NGC 容器

pytorch-quantization已预装在 NVIDIA NGC PyTorch 容器中，例如nvcr.io/nvidia/pytorch:22.12-py3，拉取对应容器镜像后可直接import pytorch_quantization使用，无需额外安装。

3. 核心概念：量化函数、描述符与量化器

理解工具包的第一步是理清“量化函数 → 量化描述符 → 量化器模块 → 量化模块”这条从底层到上层的抽象链。

3.1 两个基础量化函数：tensor_quant与fake_tensor_quant

userguide.rst 指出，pytorch_quantization/tensor_quant.py 暴露了两个基础函数：

tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False) fake_tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False)

• tensor_quant返回真正的整数量化结果（整数取值）以及缩放因子scale；
• fake_tensor_quant返回伪量化结果（仍为浮点数值，模拟量化误差但不改变数据类型），前向时数据被舍入到量化网格上。

用户指南给出的示例（固定随机种子保证可复现）：

from pytorch_quantization import tensor_quant # Generate random input. With fixed seed 12345, x should be # tensor([0.9817, 0.8796, 0.9921, 0.4611, 0.0832, 0.1784, 0.5676, 0.3376, 0.2119]) torch.manual_seed(12345) x = torch.rand(10) # fake quantize tensor x. fake_quant_x will be # tensor([0.9843, 0.8828, 0.9921, 0.4609, 0.0859, 0.1797, 0.3672, 0.5703, 0.3359, 0.2109]) fake_quant_x = tensor_quant.fake_tensor_quant(x, x.abs().max()) # quantize tensor x. quant_x will be # tensor([126., 113., 127., 59., 11., 23., 47., 73., 43., 27.]) # with scale=128.0057 quant_x, scale = tensor_quant.tensor_quant(x, x.abs().max())

从 tensor_quant.py 的_tensor_quant共享实现可以看清量化公式：

max_bound = (2.0**(num_bits - 1 + int(unsigned))) - 1.0 scale = max_bound / amax outputs = torch.clamp((inputs * scale).round_(), min_bound, max_bound)

即先由amax决定缩放因子scale = max_bound / amax，再执行round(clamp(x * scale))。实现细节还包括：

• narrow_range=True（默认）时使用对称整数区间[-127, 127]（8bit 场景），否则使用[-128, 127]；unsigned=True时区间为[0, 255]；
• 计算在 FP32 中进行以防溢出，输入/amax为 FP16 时会先提升精度；
• 当amax小于 FP16 可表示的最小值（1/(1<<24)）时，对应位置量化结果置零、scale置 1，避免除零与溢出；
• 两个函数的前向在 GPU 上走 CUDA 扩展（src/tensor_quant_gpu.cu），非 CUDA 输入则回退到 CPU 参考实现。

两者的反向传播都采用直通估计器（Straight-Through Estimator, STE）：在[-amax, amax]区间内梯度原样通过，区间外梯度置零（见 tensor_quant.py 的backward实现）。这是量化感知训练得以进行的基础。

3.2QuantDescriptor：量化方式的描述符

QuantDescriptor描述一个张量应当如何被量化（tensor_quant.py），关键参数如下：

工具包预置了一批常用描述符（tensor_quant.py）：

QUANT_DESC_8BIT_PER_TENSOR = QuantDescriptor(num_bits=8) QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0)) QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0)) QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0)) QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW = QuantDescriptor(num_bits=8, axis=(0)) QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1)) QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1)) QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))

可以看到：常规卷积/线性层的权重采用per-channel（axis=(0)），而转置卷积权重由于通道维在轴 1，采用axis=(1)；激活则统一走 per-tensor 的QUANT_DESC_8BIT_PER_TENSOR。

3.3TensorQuantizer：量化张量的模块封装

TensorQuantizer（nn/modules/tensor_quantizer.py）由QuantDescriptor定义，负责在模块前向中对张量执行量化，并维护训练所需的变量与统计。其forward（tensor_quantizer.py）按顺序经过多个阶段：

1. ONNX 导出就绪检查：当_enable_onnx_export为真时，断言量化器已完成校准（注册了_amaxbuffer），否则抛出异常；
2. pre_quant_scale：若设置了该缩放（SmoothQuant 风格的激活预缩放），先对输入乘以缩放系数；
3. disabled 旁路：disable()后整个模块直接透传输入；
4. calib 阶段：enable_calib()后调用校准器的collect(inputs)收集统计信息（此时应配合disable_quant()使用全精度数据）；
5. clip 阶段：learn_amax=True时先经可学习Clip模块裁剪，amax 由max(-clip_min, clip_max)决定；
6. quant 阶段：调用fake_tensor_quant/tensor_quant（位宽为元组(4,3)时走scaled_e4m3FP8 路径）完成量化。

配套的关键方法与属性包括：

• load_calib_amax(*args, **kwargs)（tensor_quantizer.py）：从校准器取出 amax 写入 buffer；strict=False时允许校准器未收集到数据（amax 置 NaN 并告警）；
• enable_calib/disable_calib/enable_quant/disable_quant/enable/disable/enable_clip/disable_clip：各阶段的开关；
• step_size属性：amax / (2^(num_bits-1+unsigned) - 1)，即量化步长；E4M3 时maxbound=448.0。

3.4 量化模块家族

pytorch_quantization.nn（常以quant_nn导入）提供了与torch.nn对应的量化版本，例如QuantConv1d/2d/3d、QuantConvTranspose*、QuantLinear、QuantLSTM/LSTMCell、QuantAvgPool*、QuantAdaptiveAvgPool*等（完整替换映射见 quant_modules.py）。这些模块在原有参数基础上额外接收quant_desc_input与quant_desc_weight两个描述符，并设置了合理的默认值——以 nn/modules/quant_conv.py 与 nn/modules/quant_linear.py 为例：

• QuantConv2d：输入默认QUANT_DESC_8BIT_PER_TENSOR，权重默认QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL；
• QuantLinear：输入默认 per-tensor，权重默认QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW（per-row/per-channel）。

手动构造量化模块的示例（摘自 userguide.rst）：

from torch import nn from pytorch_quantization import tensor_quant import pytorch_quantization.nn as quant_nn # pytorch's module fc1 = nn.Linear(in_features, out_features, bias=True) conv1 = nn.Conv2d(in_channels, out_channels, kernel_size) # quantized version quant_fc1 = quant_nn.Linear( in_features, out_features, bias=True, quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR, quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW) quant_conv1 = quant_nn.Conv2d( in_channels, out_channels, kernel_size, quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR, quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)

4. 自动量化与手动量化

4.1 自动量化：quant_modules.initialize()

对于完全由标准模块组成的模型，后训练量化（PTQ）只需一行调用（userguide.rst）：

from pytorch_quantization import quant_modules model = torchvision.models.resnet50()

其底层机制是monkey patching（见 quant_modules.py）：initialize()通过全局辅助对象QuantModuleReplacementHelper将torch.nn中的Conv2d、Linear、LSTM、AvgPool2d等（见_DEFAULT_QUANT_MAP）替换为对应的quant_nn量化版本，并保存原始模块以便恢复。两个可选参数提供了细粒度控制：

• float_module_list：拒绝列表，其中的模块名不做替换（保持浮点）；
• custom_quant_modules：自定义映射，用于为torch.nn之外的模块指定量化版本。

配套的quant_modules.deactivate()会把被替换的模块恢复为原始浮点版本；enable_onnx_export()是上下文管理器，进入后设置TensorQuantizer._enable_onnx_export = True，用于 ONNX 导出前开启就绪检查（见 quant_modules.py）。

4.2 手动量化

如果模型并非完全由模块构成（例如某些激活路径是裸函数），则需要手动创建TensorQuantizer并插入到模型中的正确位置。用户指南给出了组合使用的示例：

from pytorch_quantization.tensor_quant import QuantDescriptor from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer quant_desc = QuantDescriptor(num_bits=4, fake_quant=False, axis=(0), unsigned=True) quantizer = TensorQuantizer(quant_desc) torch.manual_seed(12345) x = torch.rand(10, 9, 8, 7) quant_x = quantizer(x)

注意：QuantDescriptor的axis语义与max()的 axis相反，它指定的是“保留并分别缩放”的轴；若描述符中给出amax则直接使用，否则按axis实时计算 amax 后再量化。自定义量化模块的完整范例可参考 docs/source/tutorials/creating_custom_quantized_modules.rst。

5. 校准（Calibration）：PTQ 的关键步骤

校准是 TensorRT 语境下的术语：向量化器喂入数据样本，为激活确定最优amax。校准得到的 amax 将作为后续量化的缩放依据，因此校准数据的选择与校准方法的选取直接决定 PTQ 精度。

5.1 四种校准方法

用户指南列出了 4 种校准方法，其中 3 种由直方图校准器实现：

底层设计上，_Calibrator是抽象基类（calib/calibrator.py），约定collect(x)（收集统计）、reset()（复位）、compute_amax()（计算 amax）三个接口：

• MaxCalibrator维护全局运行最大值，collect时对输入取绝对值并按axis归约，逐批与已收集值取max；
• HistogramCalibrator以直方图形式收集一次数据（默认num_bins=2048，仅支持 per-tensor），compute_amax(method, *, stride=1, start_bin=128, percentile=99.99)按方法计算最优阈值；stride/start_bin用于 KL 搜索网格，percentile默认 99.99。直方图既可用torch.histc在 GPU 上采集（默认torch_hist=True），也支持 numpy 采集。

5.2 校准实操代码

以下是在 ResNet50 上执行mse校准的完整流程（摘自 userguide.rst）：

# Find the TensorQuantizer and enable calibration for name, module in model.named_modules(): if name.endswith('_quantizer'): module.enable_calib() module.disable_quant() # Use full precision data to calibrate # Feeding data samples model(x) # ... # Finalize calibration for name, module in model.named_modules(): if name.endswith('_quantizer'): module.load_calib_amax() module.disable_calib() module.enable_quant() # If running on GPU, it needs to call .cuda() again because new tensors will be created by calibration process model.cuda() # Keep running the quantized model

要点总结：

• 校准阶段先enable_calib()+disable_quant()，让模型以全精度运行并让每个量化器收集统计；
• 喂入若干批校准数据（覆盖有代表性的输入分布）；
• 最后load_calib_amax()将 amax 写入 buffer，恢复disable_calib()+enable_quant()；
• 由于校准过程创建了新张量，GPU 上运行需再次调用model.cuda()；
• 校准必须在导出 ONNX 之前完成。

6. 量化感知训练（QAT）

6.1 原理

QAT 建立在 STE 导数近似之上（用户指南特别指出，由于 STE 近似使训练“感知不到”量化，作者更倾向称之为“基于 STE 的训练”而非“量化感知训练”）。前向使用伪量化（保留量化误差），反向按 STE 规则回传梯度，从而让网络在存在量化误差的条件下继续收敛。

6.2 训练策略

用户指南给出的默认 QAT 策略是：校准完成后，继续用原有训练计划的一小部分进行微调——通常约为原始训练计划的 10%，起始学习率为原始训练的 1%，采用余弦退火（cosine annealing）沿余弦周期下降的一半衰减，最终降到微调起始学习率的 1%（即原始训练学习率的 0.01%）。

6.3 经验建议

用户指南总结了基于实践的两条建议：

• 使用较小的学习率：STE 近似会引入方差，过大的学习率容易放大该方差并毁掉已训练好的网络；
• 训练期间不要（频繁）改变量化表示（scale）：每一步都改变 scale 等价于每一步都更换数据格式（如 e8m7、e5m10、e3m4 等），极易影响收敛。

6.4 训练辅助工具

optim/helper.py 提供了一批面向 QAT 的实用函数：

• match_parameters(model, patterns)：按正则模式匹配并产出参数，便于分组处理；
• group_parameters(model, patterns_list, lrs=None, momentums=None, weight_decays=None)：按模式分组生成符合 PyTorch optimizer 规范的参数组（{'params': ..., 'lr': ...}）；
• freeze_parameters(model, patterns)：按模式将参数requires_grad置 False；
• quant_weight_inplace(model)：对含_weight_quantizer的模块就地执行权重量化，可用于实现 relax-and-round 等离散优化方法（真实量化就地执行仅建议推理阶段使用）。

此外，learn_amax=True时 amax 会作为可学习参数通过Clip模块在训练中更新，TensorQuantizer.init_learn_amax()可用校准后的 amax 初始化可学习范围。

7. 导出 ONNX 并部署 TensorRT

7.1 导出目标与约束

工具包导出 ONNX 的目标是部署到 TensorRT，而非 ONNX Runtime，因此只导出 TensorRT 可消费的形态。约束条件（见 userguide.rst 的 Note）：

• 目前仅支持导出int8 与 fp8 的伪量化模块；
• 量化模块必须先完成校准再导出；
• 需要enable_onnx_checker=False（原因见下文 opset 说明）。

7.2 导出代码

标准导出流程（int8 场景，摘自 userguide）：

import pytorch_quantization from pytorch_quantization import nn as quant_nn from pytorch_quantization import quant_modules quant_modules.initialize() model = torchvision.models.resnet50() # load the calibrated model state_dict = torch.load("quant_resnet50-entropy-1024.pth", map_location="cpu") model.load_state_dict(state_dict) model.cuda() dummy_input = torch.randn(128, 3, 224, 224, device='cuda') input_names = [ "actual_input_1" ] output_names = [ "output1" ] with pytorch_quantization.enable_onnx_export(): # enable_onnx_checker needs to be disabled. See notes below. torch.onnx.export( model, dummy_input, "quant_resnet50.onnx", verbose=True, opset_version=10, enable_onnx_checker=False )

7.3 导出的底层形态

从 tensor_quant.py 的symbolic实现可以看清导出形态：

• int8 伪量化会被拆解为一对QuantizeLinear/DequantizeLinearONNX 算子（_onnx_int8_helper），TensorRT 拿到该图后会以最优化方式在 INT8 下执行；
• 用户指南 Note 提示：axis从 opset 13 起才加入QuantizeLinear/DequantizeLinear，这是示例使用opset_version=10并关闭 onnx checker 的原因之一；
• **FP8（E4M3）**导出使用自定义算子trt::TRT_FP8QuantizeLinear/trt::TRT_FP8DequantizeLinear（见_onnx_fp8_quantize/_onnx_fp8_dequantize），且 FP8 ONNX 导出要求标量scale，因此 E4M3 仅支持 per-tensor 量化（TensorQuantizer._check_onnx_readiness会对非标量 amax 断言失败）。

导出前的就绪检查逻辑位于 nn/modules/tensor_quantizer.py：未校准（无_amax）会直接断言报错；处于校准模式会给出警告；E4M3 场景强制 per-tensor amax。

8. 端到端实战：基于 classification_flow.py 的完整流水线

仓库在 examples/torchvision/classification_flow.py 提供了一个覆盖“校准 → 微调 → ONNX/TRT 导出评估”全流程的分类任务脚本，并配套两个 Notebook 教程：calibrate_quant_resnet50.ipynb（校准）与 finetune_quant_resnet50.ipynb（微调）。ResNet50 专用流程文档见 docs/source/tutorials/quant_resnet50.rst。

8.1 命令行参数总览

脚本的完整参数（classification_flow.py）：

8.2 执行流程

main()（classification_flow.py）按以下顺序执行：

1. 准备模型与数据：prepare_model()中若disable_pcq为真，则显式将 Conv/Linear 的输入与权重描述符设为 per-tensor（源码注释说明：per-channel 权重量化当前不被 ONNX-RT/PyTorch 导出支持）；模型构建优先使用models/classification/resnet.py中带quantize=True的预量化定义，否则回退到quant_modules.initialize()+ torchvision 模型 +deactivate()；
2. 初始精度评估：记录浮点基线 Top1；
3. 校准：calibrate_model()（classification_flow.py）内部调用collect_stats()与compute_amax()——collect_stats遍历TensorQuantizer打开校准、关闭量化并喂数据；compute_amax按校准器类型调用load_calib_amax(method=...)，并将带 amax 的 state_dict 保存为${model}-{method}-{num_calib_batch*batch_size}.pth（如resnet50-mse-512.pth，其中 512 = 4 批 × 128）；
4. 逐层敏感度分析（可选--sensitivity）：build_sensitivity_profile()逐层开启量化、其余层禁用，评估单层量化对精度的独立影响；
5. 微调（可选）：SGD 优化器、初始 lr=0.0001，配合CosineAnnealingLR，训练num_finetune_epochs轮；
6. 导出与评估：export_onnx导出model.onnx并（可选）用 ONNX Runtime 评估；export_trt导出model.trt并（可选）评估 TensorRT；
7. 汇总输出：以表格打印Initial / Calibrated / Finetuned / ONNX / TRT各阶段 Top1，并用--threshold校验精度是否达标。

9. 测试与验证

仓库自带的 tests/ 覆盖了工具包的几乎全部功能面，是验证理解与二次开发的绝佳参考：

• 底层量化：tensor_quant_test.py、functional_test.py 覆盖量化函数与数值正确性；
• 模块层：quant_conv_test.py、quant_conv_transposed_test.py、quant_linear_test.py、quant_pooling_test.py、quant_rnn_test.py、quant_instancenorm_test.py、quant_modules_test.py 验证各量化模块与自动替换；
• 量化器与校准：tensor_quantizer_test.py、calibrator_test.py、clip_test.py；
• 端到端：test_onnx_export.py 验证 ONNX 导出形态，integration_test.py 与 classification_flow_test.py 验证整体流程。

10. 注意事项与常见问题

1. 校准先于导出：ONNX 导出要求所有量化器均已校准（_check_onnx_readiness直接断言），未校准模型导出会报错；
2. int8/fp8 支持范围：导出仅支持 int8 与 fp8（E4M3）伪量化；E4M3 仅支持 per-tensor 量化（导出要求标量 amax）；
3. opset 与 checker：示例使用opset_version=10并关闭enable_onnx_checker，因为axis到 opset 13 才加入QuantizeLinear/DequantizeLinear；
4. per-channel 权重的部署边界：per-channel 权重量化在 PyTorch 侧训练/推理可用，但 ONNX-RT 导出评估不受支持（见classification_flow.py注释），面向 TensorRT 部署时以 TensorRT 导入结果为准；
5. 校准方法选择：max最保守，entropy（KL）与mse通常能更好权衡离群点，percentile适合存在明显长尾的场景；可通过classification_flow.py的--sensitivity逐层定位敏感层；
6. 版本演进：工具包已进入维护冻结状态，新项目建议关注 TensorRT Model Optimizer，以获取持续更新的量化与压缩能力。

至此，从安装、核心概念、PTQ 校准、QAT 微调，到 ONNX 导出并被 TensorRT 8.0+ 导入的完整链路均已覆盖；仓库内的源码与测试为每一步提供了可复现的实现依据，读者可在此基础上针对自身模型开展量化实践。

【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT