from https://leimao.github.io/blog/PyTorch-Eager-Mode-Quantization-TensorRT-Acceleration/

从 PyTorch 2.3.0 开始，PyTorch 提供了三种量化接口：eager mode 量化、FX graph mode 量化以及 PyTorch 2 Export 量化。

由于最新的 PyTorch 2 Export 量化接口阻止了量化后的 PyTorch 模型导出为 ONNX，因此若不开发自定义的 PyTorch FX graph 量化后端（比如fx2trt^[1] ），就无法使用 TensorRT 加速模型推理。

而 eager mode 量化和 FX graph mode 量化接口都支持将量化后的 PyTorch 模型导出为 ONNX，可以进一步使用 TensorRT 进行优化和加速。尽管 FX graph mode 量化接口更加灵活和强大，但某些使用场景下，使用 eager mode 量化接口仍是不可避免的。

在这篇文章中，我将展示如何使用 TensorRT 加速 PyTorch eager mode 量化接口生成的量化模型。同样的方法也适用于 FX graph mode 量化接口生成的量化模型，因为这两种量化模型都可以导出为 ONNX。

PyTorch Eager Mode 量化 TensorRT 加速

从 PyTorch eager mode 量化接口生成的量化模型使用 TensorRT 加速的过程涉及以下三个步骤：

• 在 PyTorch 中对浮点 PyTorch 模型执行 eager mode 量化，并将量化后的 PyTorch 模型导出为 ONNX。
• 修复量化后的 ONNX 模型图，以便 TensorRT 解析器能够解析它。
• 将量化后的 ONNX 模型构建为 TensorRT 引擎，进行性能分析并验证准确性。

本篇文章的源代码可以在 GitHub^[2] 上找到。

TensorRT INT8 量化要求

• TensorRT INT8 显式量化[3]要求量化模型中的权重使用每通道对称量化，激活使用每张量对称量化（per-channel symmetric quantization for weights and per-tensor symmetric quantization for activations）。因此，在 PyTorch 中进行后训练静态量化校准或量化感知训练时，必须确保量化配置满足 TensorRT INT8 量化要求。

• PyTorch Eager Mode 量化与我之前发布的 PyTorch 静态量化教程[4] 略有不同，此次的量化方法使用每通道对称量化权重和每张量对称量化激活，并将 PyTorch 的量化后端设置为 qnnpack，而不是 fbgemm。原因在于 fbgemm 不太支持 INT8 对称量化推理，这样会阻止通过跟踪将模型导出为 ONNX。

torch.backends.quantized.engine = 'qnnpack'  

per_tensor_activation_observer_range_neg_128_to_127 = torch.ao.quantization.MinMaxObserver.with_args(  
    dtype=torch.qint8,  
    qscheme=torch.per_tensor_symmetric,  
    quant_min=-128,  
    quant_max=127,  
    eps=2**-12)  
per_channel_weight_observer_range_neg_128_to_127 = torch.ao.quantization.PerChannelMinMaxObserver.with_args(  
    dtype=torch.qint8,  
    qscheme=torch.per_channel_symmetric,  
    quant_min=-128,  
    quant_max=127,  
    eps=2**-12)  
quantization_config = torch.ao.quantization.QConfig(  
    activation=per_tensor_activation_observer_range_neg_128_to_127,  
    weight=per_channel_weight_observer_range_neg_128_to_127)  

quantized_model.qconfig = quantization_config  

量化后的 ONNX 模型图修复

从 PyTorch eager mode 量化导出的量化 ONNX 模型中存在一些错误和问题，导致 TensorRT 解析器无法正确解析。因此，我们需要在构建 TensorRT 引擎之前修复量化后的 ONNX 模型。

具体来说，在量化 ONNX 模型图中会插入一个 Cast 节点，位于 QuantizeLinear 和 DequantizeLinear 节点之间，其数据类型为 uint8，而非 int8，尽管我们已显式将激活量化配置设置为 torch.qint8。因此，需要移除这些错误的 Cast 节点。

此外，Conv 节点的浮点偏置项无法被 TensorRT 解析器解析。需要计算并将其作为常量张量添加到量化后的 ONNX 模型图中。

在修复之前，量化后的 ONNX 模型图如下所示：

修复之后，量化后的 ONNX 模型图如下所示：

在某些地方，量化后的 ONNX 模型图仍可以进一步优化以获得最佳 TensorRT 性能，例如通过将跳跃连接的 Add 与之前的卷积层融合，移除 Conv 和 Add 节点之间的 QuantizeLinear 和 DequantizeLinear。

由于使用 PyTorch eager mode 量化实现起来比较棘手，本文未涉及此内容。若需获得最佳的 TensorRT 量化引擎性能，请参考 TensorRT Q/DQ 放置建议[5]，并考虑使用 NVIDIA PyTorch 量化工具包[6]。

构建和验证量化 TensorRT 引擎

浮点型 PyTorch ResNet18 模型和 INT8 量化 PyTorch ResNet18 模型在 CIFAR10 测试数据集上的准确率分别为 0.854 和 0.852。量化 TensorRT 引擎在 CIFAR10 测试数据集上的准确率为 0.851，与量化后的 PyTorch 模型一致。

对于批大小为 1、输入图像尺寸为 32 x 32 的情况下，FP16 和 INT8 ResNet18 TensorRT 引擎的推理延迟分别为 0.208177 ms 和 0.17584 ms。尽管由于小批大小和小图像尺寸导致数学利用率较低，INT8 量化 ResNet18 引擎仍比 FP16 ResNet18 引擎快 1.2 倍。如果批大小和图像尺寸增大，延迟改进将更加显著。

参考资料

• PyTorch Eager Mode Quantization TensorRT Acceleration – GitHub^[3]

• PyTorch Quantization^[4]

• PyTorch Static Quantization^[5]

公众号后台回复“数据集”获取100+深度学习各方向资源整理