PyTorch Eager 模式量化的 TensorRT 加速
内容提要
本文介绍了使用TensorRT加速PyTorch量化模型的方法，包括量化模型导出为ONNX格式，修复ONNX模型图以适应TensorRT解析器，并构建为TensorRT引擎。对比了FP16和INT8 ResNet18 TensorRT引擎的推理延迟和准确性。
延伸解读
量化接口选择：为何仍需急切模式
PyTorch 2.3.0 提供三种量化接口，但最新的 PyTorch 2 导出量化无法将模型导出为 ONNX，因此无法直接利用 TensorRT 加速。急切模式和 FX 图模式量化则支持 ONNX 导出，其中 FX 图模式更灵活，但某些场景下仍需使用急切模式。本文以急切模式为例，其方法同样适用于 FX 图模式。
TensorRT INT8 量化的硬性要求
TensorRT INT8 显式量化要求权重使用每通道对称量化，激活使用每张量对称量化。在 PyTorch 中执行训练后静态量化或量化感知训练时，必须确保量化配置满足这些要求。本文示例中，激活使用 per_tensor_symmetric，权重使用 per_channel_symmetric，并将量化后端设为 qnnpack，因为 fbgemm 对 INT8 对称量化推理支持不佳，会阻碍 ONNX 导出。
ONNX 图修复：必要步骤与注意事项
从 PyTorch 急切模式量化导出的 ONNX 模型图存在缺陷，需修复后才能被 TensorRT 解析器处理。主要问题包括：在 QuantizeLinear 和 DequantizeLinear 之间插入了错误的 Cast 节点（数据类型为 uint8 而非 int8），以及 Conv 节点的浮点偏置项无法被解析。修复方法是删除错误的 Cast 节点，并将偏置项作为常量张量添加到图中。此外，残差块跳跃连接中重复的 DequantizeLinear 节点虽可优化，但不影响正确性，可暂时忽略。
性能与精度权衡：INT8 引擎的实际收益
在 CIFAR10 测试集上，量化 TensorRT 引擎的准确率为 0.851，与量化 PyTorch 模型的 0.852 基本一致，而浮点模型为 0.854。对于批大小 1、输入 32x32 的 ResNet18，FP16 和 INT8 TensorRT 引擎的推理延迟分别为 0.208177 ms 和 0.17584 ms，INT8 实现了 1.2 倍的延迟改进。由于批大小和图像尺寸较小，计算利用率不高，若增大批大小和图像尺寸，延迟改进可能更显著。
Q&A
如何使用TensorRT加速PyTorch急切模式量化模型？
使用TensorRT加速PyTorch急切模式量化模型的步骤包括：量化模型并导出为ONNX，修复ONNX模型图，构建TensorRT引擎。
PyTorch 2.3.0支持哪些量化接口？
PyTorch 2.3.0提供急切模式量化、FX图模式量化和PyTorch 2导出量化三种接口。
TensorRT INT8量化的要求是什么？
TensorRT INT8显式量化要求权重使用每通道对称量化，激活使用每张量对称量化。
导出的量化ONNX模型图需要修复哪些问题？
导出的量化ONNX模型图需要修复错误的Cast节点和添加浮点偏置项，以便TensorRT解析器能够解析。
FP16和INT8 ResNet18 TensorRT引擎的推理延迟分别是多少？
FP16和INT8 ResNet18 TensorRT引擎的推理延迟分别为0.208177 ms和0.17584 ms，INT8引擎比FP16引擎有1.2倍的延迟改进。
量化TensorRT引擎在CIFAR10测试集上的准确率是多少？
量化TensorRT引擎在CIFAR10测试集上的准确率为0.851，与量化PyTorch模型一致。
