pytorch_quantization.nn 量化模块深度指南TensorQuantizer 与 Quant 系列层在 TensorRT 量化流程中的原理与实践【免费下载链接】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本指南以 TensorRT 开源仓库中 pytorch-quantization 子项目的 API 参考文档 nn.rst 为核心骨架系统讲解pytorch_quantization.nn这一量化神经网络模块库从一切量化逻辑的枢纽TensorQuantizer到卷积、全连接、池化、Clip、循环网络等一整套 Quant 前缀算子并深入对应的 Python 源码tensor_quantizer.py、quant_conv.py 等揭示其底层实现。读完本文你将掌握每个量化层的构造参数、默认量化描述符、前向执行流程与开关控制方法并能在自己的 PyTorch 模型中直接替换使用这些量化算子为后续 TensorRT 部署中的 PTQ训练后量化与 QAT量化感知训练流程打下基础。一、模块总览pytorch_quantization.nn提供了什么pytorch_quantization.nn是 pytorch-quantization 的量化算子库它提供的所有类被组织在两个层级中pytorch_quantization.nn.TensorQuantizer单个张量量化器模块是每个量化层内部实际执行量化动作的组件pytorch_quantization.nn.modules下的 Quant 系列层面向用户的量化网络层包括量化卷积、量化全连接、量化池化、Clip、量化 LSTM 等。从包入口 nn/init.py 可以看到整个命名空间通过通配导入组装而成from pytorch_quantization.nn.modules.tensor_quantizer import * from pytorch_quantization.nn.modules.quant_conv import * from pytorch_quantization.nn.modules.quant_linear import * from pytorch_quantization.nn.modules.quant_pooling import * from pytorch_quantization.nn.modules.clip import * from pytorch_quantization.nn.modules.quant_rnn import * from pytorch_quantization.nn.modules.quant_instancenorm import *即pytorch_quantization.nn顶层命名空间直接暴露TensorQuantizer、QuantConv1d/2d/3d、QuantConvTranspose1d/2d/3d、QuantLinear、QuantMaxPool1d/2d/3d、QuantAvgPool1d/2d/3d、QuantAdaptiveAvgPool1d/2d/3d、Clip、QuantLSTM、QuantLSTMCell以及部分未在 nn.rst 中展开但同样被导出的模块如 quant_instancenorm 中的层。值得注意的设计细节是量化层的本名与 PyTorch 原生层同名如Conv2d QuantConv2d并额外提供带 Quant 前缀的别名。源码注释明确说明了这一意图——保持与原版相同的模块名使其可以无缝替换进既有模型并加载预训练权重而从零新建模型时则鼓励显式使用 Quant 前缀别名见 quant_conv.py 顶部说明。二、TensorQuantizer一切量化的核心枢纽nn.rst 把TensorQuantizer放在文档首位它正是整个量化体系的心脏。该类继承torch.nn.Module在内部调用tensor_quant/fake_tensor_quant函数对张量执行量化并封装了训练量化网络所需的变量与滑动统计。完整实现见 tensor_quantizer.py。2.1 构造参数与 QuantDescriptorTensorQuantizer的构造函数签名如下TensorQuantizer(quant_descQuantDescriptor(), disabledFalse, if_quantTrue, if_clipFalse, if_calibFalse)各参数含义参数默认值说明quant_descQuantDescriptor()量化描述符实例定义于pytorch_quantization.tensor_quant.QuantDescriptor携带 num_bits、fake_quant、axis、scale_amax、learn_amax、unsigned、narrow_range、amax、calib_method 等全部量化设置disabledFalse若为 True整个模块被旁路bypass直接返回输入if_quantTrue若为 True执行主体量化逻辑if_clipFalse若为 True量化前先做 clip 并学习 amaxif_calibFalse若为 True运行校准流程文档注释说明校准设置后续计划并入 QuantDescriptor构造函数会逐项展开quant_desc的字段显式逐个赋值而非用 dict源码注释称这是为了获得更多控制权并据此注册_amaxbuffer、在learn_amax为 True 时创建Clip模块、以及根据calib_method创建校准器if quant_desc.calib_method histogram: self._calibrator calib.HistogramCalibrator(...) elif quant_desc.calib_method max: self._calibrator calib.MaxCalibrator(...)即校准器由 QuantDescriptor 的calib_method字段决定histogram 或 max对应 calib 文档中的校准器体系。2.2 只读与可变属性从 nn.rst 的 autoclass 展开及源码 docstring 可以看到 TensorQuantizer 的属性被明确划分为两组只读属性axis量化粒度per-tensor 时为 Noneper-channel 时为对应维度fake_quant是否伪量化fake quantization即保留浮点 dtype、但数值上模拟量化误差scale真实量化非 fake时前向计算出的缩放因子源码中scale属性在 fake 模式下会打印错误日志在尚未量化任何张量时访问则打印 critical 日志step_size步长实现为amax / (2^(num_bits - 1 unsigned) - 1)maxbound量化值上界(1 (num_bits - 1 int(unsigned))) - 1特别地当num_bits (4, 3)E4M3时返回 448.0pre_quant_scaleSmoothQuant 风格的量化前缩放因子。可变属性带 setternum_bits量化位宽整数或 E4M3 场景下的(4, 3)元组unsigned是否无符号量化narrow_range是否使用窄范围如 INT8 的 [-127, 127]amax量化绝对最大值buffersetter 支持形状校验与设备迁移。2.3 前向执行流程五级流水线TensorQuantizer.forward是理解整个量化行为的关键其执行顺序见 tensor_quantizer.pyforward方法可以归纳为如下流水线ONNX 导出就绪检查当类变量_enable_onnx_export为 True 时调用_check_onnx_readiness校验已校准存在_amaxbuffer、且 E4M3 场景下 amax 必须是标量per-tensorSmoothQuant 预缩放若设置了pre_quant_scale先对输入逐元素相乘inputs inputs * self.pre_quant_scale旁路检查_disabled为 True 则直接返回输入校准收集_if_calib为 True 时把输入交给校准器collect(inputs)校准器在见到第一个张量时才确定形状Clip 学习阶段_if_clip为 True 时对输入做 clip若未开启 learn_amax 则抛 RuntimeError量化主流程_if_quant为 True 时执行_quant_forward(inputs)——若num_bits是普通整数则走fake_tensor_quant/tensor_quant路径若num_bits是(E, M)元组则走scaled_e4m3(inputs, amax, E, M)浮点格式量化路径。_quant_forward内部还会区分两种 amax 来源learn_amaxTrue时从 Clip 的可学习上下界推导amax max(-clip_value_min, clip_value_max)否则通过_get_amax从 buffer 读取静态 amax或按axis动态归约计算输入的实际最大值。_get_amax还支持scale_amax缩放以及将低精度 dtype 的 amax 提升为 float32。2.4 校准与开关控制方法族TensorQuantizer 提供了一套细粒度的运行时控制方法用于编排校准 → 固定 amax → 训练/推理的完整生命周期旁路控制disable()/enable()量化控制disable_quant()/enable_quant()Clip 学习控制enable_clip()要求learn_amaxTrue否则抛 ValueError/disable_clip()校准控制enable_calib()校准器为 None 时抛 ValueError/disable_calib()load_calib_amax(*args, **kwargs)从校准器compute_amax的结果装载 amax 到 buffer校准器未见到任何张量时默认抛 RuntimeError传strictFalse则降级为警告并把 amax 置为 NaNinit_learn_amax()以固定 amax 作为可学习 amax 的初值per-channel 场景不支持学习取 max 初始化。此外extra_repr()会把模块当前状态序列化成一行可读字符串如8bit fake per-tensor amax1.0000 calibratorMaxCalibrator quant便于打印模型时直观确认每个量化器的配置。三、QuantMixin 与 QuantInputMixin量化层的公共底座nn.rst 中所有 Quant 系列层都建立在 modules/_utils.py 的两个 Mixin 之上QuantMixin用于既有输入又有权重的层卷积、全连接、RNN。提供类级默认描述符default_quant_desc_input/default_quant_desc_weight默认均为QUANT_DESC_8BIT_PER_TENSOR、类方法set_default_quant_desc_input/set_default_quant_desc_weight做深拷贝后写入类属性、init_quantizer与只读属性input_quantizer/weight_quantizer。init_quantizer只支持伪量化任一描述符fake_quantFalse都会抛 ValueError并支持num_layers参数批量创建nn.ModuleList形式的量化器列表RNN 使用。QuantInputMixin用于只需量化输入的单输入层池化。只维护default_quant_desc_input与单个_input_quantizer。配合pop_quant_desc_in_kwargs(quant_cls, input_onlyFalse, **kwargs)各量化层的构造函数可以从**kwargs中弹出quant_desc_input/quant_desc_weight未显式传入时回退到类默认值若有剩余未使用的 kwarg 则抛出TypeError(Unused keys: ...)保证参数使用的严格性。四、量化卷积层QuantConvNd 家族nn.rst 将_QuantConvNd与六个具体卷积层并列展开QuantConv1d、QuantConv2d、QuantConv3d、QuantConvTranspose1d、QuantConvTranspose2d、QuantConvTranspose3d实现位于 quant_conv.py。4.1 基类设计_QuantConvNd同时继承torch.nn.modules.conv._ConvNd与_utils.QuantMixin基类默认描述符为QUANT_DESC_8BIT_PER_TENSOR构造时除标准卷积参数外还接收quant_desc_input/quant_desc_weight。其_quant(input)方法完成了量化层的核心约定——对输入和权重分别套用量化器def _quant(self, input): quant_input self._input_quantizer(input) quant_weight self._weight_quantizer(self.weight) return (quant_input, quant_weight)4.2 各卷积层的默认权重描述符每个具体层覆盖了更适合自身结构的默认权重量化描述符输入仍为 per-tensor 8bit层默认权重描述符量化粒度QuantConv1dQUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL权重 per-channelQuantConv2dQUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL权重 per-channelQuantConv3dQUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL权重 per-channelQuantConvTranspose1dQUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL权重 per-channelQuantConvTranspose2dQUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL权重 per-channelQuantConvTranspose3dQUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL权重 per-channel这意味着卷积层在未显式指定描述符时输入采用 per-tensor 8bit 量化权重采用 per-channel 8bit 量化——这正是 INT8 卷积在 TensorRT 等推理引擎中推荐的量化配置。4.3 前向实现细节各具体层在forward中先调用_quant得到量化后的输入与权重再交给torch.nn.functional完成计算。以QuantConv2d为例标准路径直接调用F.conv2d(quant_input, quant_weight, self.bias, ...)当padding_mode circular时则先手动展开 padding 再卷积。QuantConvTranspose*额外要求padding_mode zeros否则抛 ValueError并使用_output_padding_nd兼容不同 PyTorch 版本中_output_padding的签名差异随后调用F.conv_transpose*。五、QuantLinear全连接层的量化QuantLinear 继承nn.Linear与QuantMixin语义为y dequant(quant(x) · quant(A)^T b)。其默认描述符为输入QUANT_DESC_8BIT_PER_TENSORper-tensor 8bit权重QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW按行 per-row 8bit即按输出神经元维度量化权重。构造函数签名QuantLinear(in_features, out_features, biasTrue, **kwargs)量化描述符从 kwargs 弹出。前向实现非常简洁def forward(self, input): quant_input self._input_quantizer(input) quant_weight self._weight_quantizer(self.weight) output F.linear(quant_input, quant_weight, biasself.bias) return output同样模块保留Linear本名并定义Linear QuantLinear别名。六、量化池化层输入侧量化nn.rst 展开的池化层共九种QuantMaxPool1d/2d/3d、QuantAvgPool1d/2d/3d、QuantAdaptiveAvgPool1d/2d/3d实现在 quant_pooling.py。这些层继承_utils.QuantInputMixin因此只量化输入、不量化权重池化本身无权重。各层构造函数与 PyTorch 原生池化签名一致如QuantMaxPool2d(kernel_size, strideNone, padding0, dilation1, return_indicesFalse, ceil_modeFalse, **kwargs)从 kwargs 中弹出quant_desc_input后前向先对输入套用_input_quantizer再调用父类原生池化def forward(self, input): quant_input self._input_quantizer(input) return super(QuantMaxPool2d, self).forward(quant_input)QuantAdaptiveAvgPool*的签名只有output_size一个必选参数。文件末尾同样定义无 Quant 前缀别名MaxPool2d QuantMaxPool2d等。七、Clip可学习量化范围的基石Clip模块实现在 clip.py是 learn_amax 模式量化感知训练中的范围学习的关键组件。PyTorch 原生只有简单的 clamp 函数而Clip将上下界建模为可学习参数Clip(clip_value_min, clip_value_max, learn_minFalse, learn_maxFalse)clip_value_min/clip_value_max下界与上界标量或张量learn_min/learn_max为 True 时对应界被包装为torch.nn.parameter.Parameter参与梯度更新并用传入值初始化为 False 时直接保存原值。源码还约束当开启 learn 时初始值必须是标量非标量会抛 ValueError。前向调用QF.clip(inputs, min, max)见 functional.rst 对应的函数式接口。TensorQuantizer 在learn_amaxTrue时自动创建一个Clip(-init_amax, init_amax, learn_minTrue, learn_maxTrue)这就是先学范围、再启用量化的实验性 clip 阶段的实现载体。八、量化循环网络QuantLSTM 与 QuantLSTMCellnn.rst 收尾处展开的QuantLSTM、QuantLSTMCell实现在 quant_rnn.py代码源自 PyTorch 早期版本的 RNN 实现后端替换为pytorch_quantization.nn._functions.quant_rnn。8.1 QuantLSTM基于 QuantRNNBaseQuantLSTM(input_size, hidden_size, num_layers1, biasTrue, batch_firstFalse, dropout0, bidirectionalFalse, proj_size0, **kwargs)继承QuantRNNBase。要点默认权重描述符为QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW与 QuantLinear 一致proj_size 0时抛出 proj_size is not supported in pytorch-quantization yet即暂不支持投影 LSTM内部按num_layers × (1 bidirectional)数量通过init_quantizer(..., num_layers...)创建_input_quantizers/_weight_quantizers两个nn.ModuleListforward支持PackedSequence变长输入最终调用quant_rnn.RNN(...)函数并把每层的量化器列表一并传入func(input, self.all_weights, hx, batch_sizes, self._input_quantizers, self._weight_quantizers)——即对每个门控矩阵weight_ih/weight_hh与输入分别量化后再执行循环计算保留LSTM QuantLSTM别名。8.2 QuantLSTMCellQuantLSTMCell(input_size, hidden_size, biasTrue, **kwargs)实现单步 LSTM 单元权重形状为4 * hidden_size四门合并默认权重描述符同样为 per-row 8bit。前向在完成输入/隐藏状态校验后调用quant_rnn.LSTMCell(input, hx, weight_ih, weight_hh, bias_ih, bias_hh, self._input_quantizer, self._weight_quantizer)把量化器直接传入底层函数执行量化后的单元计算。该文件还包含未加 Quant 前缀的QuantRNN、GRU、QuantRNNCell、GRUCell等实现部分未展开量化器属于辅助/兼容实现。九、实战接入把 Quant 层用进你的模型9.1 直接构建量化网络从pytorch_quantization.nn直接导入并替换对应层即可import torch.nn as nn from pytorch_quantization import tensor_quant from pytorch_quantization.nn import QuantConv2d, QuantLinear, TensorQuantizer class QuantNet(nn.Module): def __init__(self): super().__init__() self.conv QuantConv2d(3, 64, kernel_size3, padding1) self.fc QuantLinear(64 * 8 * 8, 10) # 也可通过 kwargs 显式指定量化描述符 # self.conv QuantConv2d(3, 64, 3, padding1, # quant_desc_inputtensor_quant.QUANT_DESC_8BIT_PER_TENSOR, # quant_desc_weighttensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL) def forward(self, x): x self.conv(x) x x.flatten(1) return self.fc(x)9.2 自动替换既有模型仓库还提供pytorch_quantization.quant_modules见 quant_modules.py可在不改动模型源码的前提下通过quant_modules.initialize()把nn.Conv2d等原生层自动替换为对应 Quant 层配合 quant_resnet50.rst 教程实现 ResNet-50 的完整量化流程。更详细的用法说明见 userguide.rst 与 creating_custom_quantized_modules.rst后者讲解如何基于 QuantMixin 自定义量化模块。9.3 量化的完整生命周期结合 calib.rst 与 optim.rst一个典型的量化流程是先用校准数据集在_if_calibTrue模式下让各层的 TensorQuantizer 收集激活分布 →load_calib_amax()装载校准得到的 amax → 固定 amax 后进行量化感知训练learn_amax 模式经Clip学习范围→ 最后导出。量化算子与底层数学tensor_quant、fake_tensor_quant、QuantDescriptor及各类QUANT_DESC_*默认描述符的细节可查阅 tensor_quant.rst这些默认描述符正是上文各层default_quant_desc_*类属性的取值来源。十、总结与文档索引pytorch_quantization.nn以TensorQuantizer为执行核心、以QuantMixin/QuantInputMixin为架构底座构建了一套与 PyTorch 原生模块一一对应、可无缝替换的量化算子族双通道量化带权重的层同时量化输入与权重且默认采用输入 per-tensor、权重 per-channel/per-row的 8bit 配置无权重层池化只量化输入三阶段生命周期校准collect → load_calib_amax、固定 amax 推理/训练、learn_amax 范围学习Clip由 TensorQuantizer 的开关方法族精细控制扩展能力E4M3num_bits(4,3)浮点格式量化、SmoothQuant 的pre_quant_scale、ONNX 导出就绪检查等实验性特性均内置于 TensorQuantizer。相关仓库资源速查API 参考文档nn.rst、tensor_quant.rst、calib.rst、functional.rst、optim.rst使用指南与教程userguide.rst、quant_resnet50.rst、creating_custom_quantized_modules.rst核心源码tensor_quantizer.py、_utils.py、quant_conv.py、quant_linear.py、quant_pooling.py、clip.py、quant_rnn.py掌握这套量化算子后结合 TensorRT 对 INT8 引擎的构建能力即可在 PyTorch 侧完成量化建模 → 校准 → 训练/微调的上游工作为后续导出与部署铺平道路。【免费下载链接】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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
pytorch_quantization.nn 量化模块深度指南TensorQuantizer 与 Quant 系列层在 TensorRT 量化流程中的原理与实践【免费下载链接】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本指南以 TensorRT 开源仓库中 pytorch-quantization 子项目的 API 参考文档 nn.rst 为核心骨架系统讲解pytorch_quantization.nn这一量化神经网络模块库从一切量化逻辑的枢纽TensorQuantizer到卷积、全连接、池化、Clip、循环网络等一整套 Quant 前缀算子并深入对应的 Python 源码tensor_quantizer.py、quant_conv.py 等揭示其底层实现。读完本文你将掌握每个量化层的构造参数、默认量化描述符、前向执行流程与开关控制方法并能在自己的 PyTorch 模型中直接替换使用这些量化算子为后续 TensorRT 部署中的 PTQ训练后量化与 QAT量化感知训练流程打下基础。一、模块总览pytorch_quantization.nn提供了什么pytorch_quantization.nn是 pytorch-quantization 的量化算子库它提供的所有类被组织在两个层级中pytorch_quantization.nn.TensorQuantizer单个张量量化器模块是每个量化层内部实际执行量化动作的组件pytorch_quantization.nn.modules下的 Quant 系列层面向用户的量化网络层包括量化卷积、量化全连接、量化池化、Clip、量化 LSTM 等。从包入口 nn/init.py 可以看到整个命名空间通过通配导入组装而成from pytorch_quantization.nn.modules.tensor_quantizer import * from pytorch_quantization.nn.modules.quant_conv import * from pytorch_quantization.nn.modules.quant_linear import * from pytorch_quantization.nn.modules.quant_pooling import * from pytorch_quantization.nn.modules.clip import * from pytorch_quantization.nn.modules.quant_rnn import * from pytorch_quantization.nn.modules.quant_instancenorm import *即pytorch_quantization.nn顶层命名空间直接暴露TensorQuantizer、QuantConv1d/2d/3d、QuantConvTranspose1d/2d/3d、QuantLinear、QuantMaxPool1d/2d/3d、QuantAvgPool1d/2d/3d、QuantAdaptiveAvgPool1d/2d/3d、Clip、QuantLSTM、QuantLSTMCell以及部分未在 nn.rst 中展开但同样被导出的模块如 quant_instancenorm 中的层。值得注意的设计细节是量化层的本名与 PyTorch 原生层同名如Conv2d QuantConv2d并额外提供带 Quant 前缀的别名。源码注释明确说明了这一意图——保持与原版相同的模块名使其可以无缝替换进既有模型并加载预训练权重而从零新建模型时则鼓励显式使用 Quant 前缀别名见 quant_conv.py 顶部说明。二、TensorQuantizer一切量化的核心枢纽nn.rst 把TensorQuantizer放在文档首位它正是整个量化体系的心脏。该类继承torch.nn.Module在内部调用tensor_quant/fake_tensor_quant函数对张量执行量化并封装了训练量化网络所需的变量与滑动统计。完整实现见 tensor_quantizer.py。2.1 构造参数与 QuantDescriptorTensorQuantizer的构造函数签名如下TensorQuantizer(quant_descQuantDescriptor(), disabledFalse, if_quantTrue, if_clipFalse, if_calibFalse)各参数含义参数默认值说明quant_descQuantDescriptor()量化描述符实例定义于pytorch_quantization.tensor_quant.QuantDescriptor携带 num_bits、fake_quant、axis、scale_amax、learn_amax、unsigned、narrow_range、amax、calib_method 等全部量化设置disabledFalse若为 True整个模块被旁路bypass直接返回输入if_quantTrue若为 True执行主体量化逻辑if_clipFalse若为 True量化前先做 clip 并学习 amaxif_calibFalse若为 True运行校准流程文档注释说明校准设置后续计划并入 QuantDescriptor构造函数会逐项展开quant_desc的字段显式逐个赋值而非用 dict源码注释称这是为了获得更多控制权并据此注册_amaxbuffer、在learn_amax为 True 时创建Clip模块、以及根据calib_method创建校准器if quant_desc.calib_method histogram: self._calibrator calib.HistogramCalibrator(...) elif quant_desc.calib_method max: self._calibrator calib.MaxCalibrator(...)即校准器由 QuantDescriptor 的calib_method字段决定histogram 或 max对应 calib 文档中的校准器体系。2.2 只读与可变属性从 nn.rst 的 autoclass 展开及源码 docstring 可以看到 TensorQuantizer 的属性被明确划分为两组只读属性axis量化粒度per-tensor 时为 Noneper-channel 时为对应维度fake_quant是否伪量化fake quantization即保留浮点 dtype、但数值上模拟量化误差scale真实量化非 fake时前向计算出的缩放因子源码中scale属性在 fake 模式下会打印错误日志在尚未量化任何张量时访问则打印 critical 日志step_size步长实现为amax / (2^(num_bits - 1 unsigned) - 1)maxbound量化值上界(1 (num_bits - 1 int(unsigned))) - 1特别地当num_bits (4, 3)E4M3时返回 448.0pre_quant_scaleSmoothQuant 风格的量化前缩放因子。可变属性带 setternum_bits量化位宽整数或 E4M3 场景下的(4, 3)元组unsigned是否无符号量化narrow_range是否使用窄范围如 INT8 的 [-127, 127]amax量化绝对最大值buffersetter 支持形状校验与设备迁移。2.3 前向执行流程五级流水线TensorQuantizer.forward是理解整个量化行为的关键其执行顺序见 tensor_quantizer.pyforward方法可以归纳为如下流水线ONNX 导出就绪检查当类变量_enable_onnx_export为 True 时调用_check_onnx_readiness校验已校准存在_amaxbuffer、且 E4M3 场景下 amax 必须是标量per-tensorSmoothQuant 预缩放若设置了pre_quant_scale先对输入逐元素相乘inputs inputs * self.pre_quant_scale旁路检查_disabled为 True 则直接返回输入校准收集_if_calib为 True 时把输入交给校准器collect(inputs)校准器在见到第一个张量时才确定形状Clip 学习阶段_if_clip为 True 时对输入做 clip若未开启 learn_amax 则抛 RuntimeError量化主流程_if_quant为 True 时执行_quant_forward(inputs)——若num_bits是普通整数则走fake_tensor_quant/tensor_quant路径若num_bits是(E, M)元组则走scaled_e4m3(inputs, amax, E, M)浮点格式量化路径。_quant_forward内部还会区分两种 amax 来源learn_amaxTrue时从 Clip 的可学习上下界推导amax max(-clip_value_min, clip_value_max)否则通过_get_amax从 buffer 读取静态 amax或按axis动态归约计算输入的实际最大值。_get_amax还支持scale_amax缩放以及将低精度 dtype 的 amax 提升为 float32。2.4 校准与开关控制方法族TensorQuantizer 提供了一套细粒度的运行时控制方法用于编排校准 → 固定 amax → 训练/推理的完整生命周期旁路控制disable()/enable()量化控制disable_quant()/enable_quant()Clip 学习控制enable_clip()要求learn_amaxTrue否则抛 ValueError/disable_clip()校准控制enable_calib()校准器为 None 时抛 ValueError/disable_calib()load_calib_amax(*args, **kwargs)从校准器compute_amax的结果装载 amax 到 buffer校准器未见到任何张量时默认抛 RuntimeError传strictFalse则降级为警告并把 amax 置为 NaNinit_learn_amax()以固定 amax 作为可学习 amax 的初值per-channel 场景不支持学习取 max 初始化。此外extra_repr()会把模块当前状态序列化成一行可读字符串如8bit fake per-tensor amax1.0000 calibratorMaxCalibrator quant便于打印模型时直观确认每个量化器的配置。三、QuantMixin 与 QuantInputMixin量化层的公共底座nn.rst 中所有 Quant 系列层都建立在 modules/_utils.py 的两个 Mixin 之上QuantMixin用于既有输入又有权重的层卷积、全连接、RNN。提供类级默认描述符default_quant_desc_input/default_quant_desc_weight默认均为QUANT_DESC_8BIT_PER_TENSOR、类方法set_default_quant_desc_input/set_default_quant_desc_weight做深拷贝后写入类属性、init_quantizer与只读属性input_quantizer/weight_quantizer。init_quantizer只支持伪量化任一描述符fake_quantFalse都会抛 ValueError并支持num_layers参数批量创建nn.ModuleList形式的量化器列表RNN 使用。QuantInputMixin用于只需量化输入的单输入层池化。只维护default_quant_desc_input与单个_input_quantizer。配合pop_quant_desc_in_kwargs(quant_cls, input_onlyFalse, **kwargs)各量化层的构造函数可以从**kwargs中弹出quant_desc_input/quant_desc_weight未显式传入时回退到类默认值若有剩余未使用的 kwarg 则抛出TypeError(Unused keys: ...)保证参数使用的严格性。四、量化卷积层QuantConvNd 家族nn.rst 将_QuantConvNd与六个具体卷积层并列展开QuantConv1d、QuantConv2d、QuantConv3d、QuantConvTranspose1d、QuantConvTranspose2d、QuantConvTranspose3d实现位于 quant_conv.py。4.1 基类设计_QuantConvNd同时继承torch.nn.modules.conv._ConvNd与_utils.QuantMixin基类默认描述符为QUANT_DESC_8BIT_PER_TENSOR构造时除标准卷积参数外还接收quant_desc_input/quant_desc_weight。其_quant(input)方法完成了量化层的核心约定——对输入和权重分别套用量化器def _quant(self, input): quant_input self._input_quantizer(input) quant_weight self._weight_quantizer(self.weight) return (quant_input, quant_weight)4.2 各卷积层的默认权重描述符每个具体层覆盖了更适合自身结构的默认权重量化描述符输入仍为 per-tensor 8bit层默认权重描述符量化粒度QuantConv1dQUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL权重 per-channelQuantConv2dQUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL权重 per-channelQuantConv3dQUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL权重 per-channelQuantConvTranspose1dQUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL权重 per-channelQuantConvTranspose2dQUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL权重 per-channelQuantConvTranspose3dQUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL权重 per-channel这意味着卷积层在未显式指定描述符时输入采用 per-tensor 8bit 量化权重采用 per-channel 8bit 量化——这正是 INT8 卷积在 TensorRT 等推理引擎中推荐的量化配置。4.3 前向实现细节各具体层在forward中先调用_quant得到量化后的输入与权重再交给torch.nn.functional完成计算。以QuantConv2d为例标准路径直接调用F.conv2d(quant_input, quant_weight, self.bias, ...)当padding_mode circular时则先手动展开 padding 再卷积。QuantConvTranspose*额外要求padding_mode zeros否则抛 ValueError并使用_output_padding_nd兼容不同 PyTorch 版本中_output_padding的签名差异随后调用F.conv_transpose*。五、QuantLinear全连接层的量化QuantLinear 继承nn.Linear与QuantMixin语义为y dequant(quant(x) · quant(A)^T b)。其默认描述符为输入QUANT_DESC_8BIT_PER_TENSORper-tensor 8bit权重QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW按行 per-row 8bit即按输出神经元维度量化权重。构造函数签名QuantLinear(in_features, out_features, biasTrue, **kwargs)量化描述符从 kwargs 弹出。前向实现非常简洁def forward(self, input): quant_input self._input_quantizer(input) quant_weight self._weight_quantizer(self.weight) output F.linear(quant_input, quant_weight, biasself.bias) return output同样模块保留Linear本名并定义Linear QuantLinear别名。六、量化池化层输入侧量化nn.rst 展开的池化层共九种QuantMaxPool1d/2d/3d、QuantAvgPool1d/2d/3d、QuantAdaptiveAvgPool1d/2d/3d实现在 quant_pooling.py。这些层继承_utils.QuantInputMixin因此只量化输入、不量化权重池化本身无权重。各层构造函数与 PyTorch 原生池化签名一致如QuantMaxPool2d(kernel_size, strideNone, padding0, dilation1, return_indicesFalse, ceil_modeFalse, **kwargs)从 kwargs 中弹出quant_desc_input后前向先对输入套用_input_quantizer再调用父类原生池化def forward(self, input): quant_input self._input_quantizer(input) return super(QuantMaxPool2d, self).forward(quant_input)QuantAdaptiveAvgPool*的签名只有output_size一个必选参数。文件末尾同样定义无 Quant 前缀别名MaxPool2d QuantMaxPool2d等。七、Clip可学习量化范围的基石Clip模块实现在 clip.py是 learn_amax 模式量化感知训练中的范围学习的关键组件。PyTorch 原生只有简单的 clamp 函数而Clip将上下界建模为可学习参数Clip(clip_value_min, clip_value_max, learn_minFalse, learn_maxFalse)clip_value_min/clip_value_max下界与上界标量或张量learn_min/learn_max为 True 时对应界被包装为torch.nn.parameter.Parameter参与梯度更新并用传入值初始化为 False 时直接保存原值。源码还约束当开启 learn 时初始值必须是标量非标量会抛 ValueError。前向调用QF.clip(inputs, min, max)见 functional.rst 对应的函数式接口。TensorQuantizer 在learn_amaxTrue时自动创建一个Clip(-init_amax, init_amax, learn_minTrue, learn_maxTrue)这就是先学范围、再启用量化的实验性 clip 阶段的实现载体。八、量化循环网络QuantLSTM 与 QuantLSTMCellnn.rst 收尾处展开的QuantLSTM、QuantLSTMCell实现在 quant_rnn.py代码源自 PyTorch 早期版本的 RNN 实现后端替换为pytorch_quantization.nn._functions.quant_rnn。8.1 QuantLSTM基于 QuantRNNBaseQuantLSTM(input_size, hidden_size, num_layers1, biasTrue, batch_firstFalse, dropout0, bidirectionalFalse, proj_size0, **kwargs)继承QuantRNNBase。要点默认权重描述符为QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW与 QuantLinear 一致proj_size 0时抛出 proj_size is not supported in pytorch-quantization yet即暂不支持投影 LSTM内部按num_layers × (1 bidirectional)数量通过init_quantizer(..., num_layers...)创建_input_quantizers/_weight_quantizers两个nn.ModuleListforward支持PackedSequence变长输入最终调用quant_rnn.RNN(...)函数并把每层的量化器列表一并传入func(input, self.all_weights, hx, batch_sizes, self._input_quantizers, self._weight_quantizers)——即对每个门控矩阵weight_ih/weight_hh与输入分别量化后再执行循环计算保留LSTM QuantLSTM别名。8.2 QuantLSTMCellQuantLSTMCell(input_size, hidden_size, biasTrue, **kwargs)实现单步 LSTM 单元权重形状为4 * hidden_size四门合并默认权重描述符同样为 per-row 8bit。前向在完成输入/隐藏状态校验后调用quant_rnn.LSTMCell(input, hx, weight_ih, weight_hh, bias_ih, bias_hh, self._input_quantizer, self._weight_quantizer)把量化器直接传入底层函数执行量化后的单元计算。该文件还包含未加 Quant 前缀的QuantRNN、GRU、QuantRNNCell、GRUCell等实现部分未展开量化器属于辅助/兼容实现。九、实战接入把 Quant 层用进你的模型9.1 直接构建量化网络从pytorch_quantization.nn直接导入并替换对应层即可import torch.nn as nn from pytorch_quantization import tensor_quant from pytorch_quantization.nn import QuantConv2d, QuantLinear, TensorQuantizer class QuantNet(nn.Module): def __init__(self): super().__init__() self.conv QuantConv2d(3, 64, kernel_size3, padding1) self.fc QuantLinear(64 * 8 * 8, 10) # 也可通过 kwargs 显式指定量化描述符 # self.conv QuantConv2d(3, 64, 3, padding1, # quant_desc_inputtensor_quant.QUANT_DESC_8BIT_PER_TENSOR, # quant_desc_weighttensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL) def forward(self, x): x self.conv(x) x x.flatten(1) return self.fc(x)9.2 自动替换既有模型仓库还提供pytorch_quantization.quant_modules见 quant_modules.py可在不改动模型源码的前提下通过quant_modules.initialize()把nn.Conv2d等原生层自动替换为对应 Quant 层配合 quant_resnet50.rst 教程实现 ResNet-50 的完整量化流程。更详细的用法说明见 userguide.rst 与 creating_custom_quantized_modules.rst后者讲解如何基于 QuantMixin 自定义量化模块。9.3 量化的完整生命周期结合 calib.rst 与 optim.rst一个典型的量化流程是先用校准数据集在_if_calibTrue模式下让各层的 TensorQuantizer 收集激活分布 →load_calib_amax()装载校准得到的 amax → 固定 amax 后进行量化感知训练learn_amax 模式经Clip学习范围→ 最后导出。量化算子与底层数学tensor_quant、fake_tensor_quant、QuantDescriptor及各类QUANT_DESC_*默认描述符的细节可查阅 tensor_quant.rst这些默认描述符正是上文各层default_quant_desc_*类属性的取值来源。十、总结与文档索引pytorch_quantization.nn以TensorQuantizer为执行核心、以QuantMixin/QuantInputMixin为架构底座构建了一套与 PyTorch 原生模块一一对应、可无缝替换的量化算子族双通道量化带权重的层同时量化输入与权重且默认采用输入 per-tensor、权重 per-channel/per-row的 8bit 配置无权重层池化只量化输入三阶段生命周期校准collect → load_calib_amax、固定 amax 推理/训练、learn_amax 范围学习Clip由 TensorQuantizer 的开关方法族精细控制扩展能力E4M3num_bits(4,3)浮点格式量化、SmoothQuant 的pre_quant_scale、ONNX 导出就绪检查等实验性特性均内置于 TensorQuantizer。相关仓库资源速查API 参考文档nn.rst、tensor_quant.rst、calib.rst、functional.rst、optim.rst使用指南与教程userguide.rst、quant_resnet50.rst、creating_custom_quantized_modules.rst核心源码tensor_quantizer.py、_utils.py、quant_conv.py、quant_linear.py、quant_pooling.py、clip.py、quant_rnn.py掌握这套量化算子后结合 TensorRT 对 INT8 引擎的构建能力即可在 PyTorch 侧完成量化建模 → 校准 → 训练/微调的上游工作为后续导出与部署铺平道路。【免费下载链接】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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
