tensorrt官方int8量化方法汇总

原理及操作

量化的基本原理及流程可参看懂你的神经网络量化教程：第一讲、量化番外篇、TensorRT中的INT8、tensorRT int8量化示例代码

Tensorrt

方式1：trtexec（PTQ的一种）

1. int8量化

trtexec --onnx=XX.onnx --saveEngine=model.plan --int8 --workspace=4096

• 精度损失很大，不建议直接采用；
• trtexec 有提供 --calib=接口进行校正，但需要对中间特征进行cache文件保存，比较麻烦，官方文档也是采用上述方式进行int8量化；
• 与fp16的模型在测试集上测试指标，可以看到精度下降非常严重；

2. int8 fp16混合量化

trtexec --onnx=XX.onnx --saveEngine=model.plan --int8 --fp16 --workspace=4096

• 测试集上统计指标：相比纯int8量化，效果要好，但是相比fp16，精度下降依然非常严重

方式2：PTQ

1. engine序列化时执行

注：属于隐式量化

1.1 python onnx转trt

• 操作流程：按照常规方案导出onnx，onnx序列化为tensorrt engine之前打开int8量化模式并采用校正数据集进行校正；

• 优点：1. 导出onnx之前的所有操作都为常规操作；2. 相比在pytorch中进行PTQ int8量化，所需显存小；

• 缺点：1. 量化过程为黑盒子，无法看到中间过程；2. 校正过程需在实际运行的tensorrt版本中进行并保存tensorrt engine；3.量化过程中发现，即使模型为动态输入，校正数据集使用时也必须与推理时的输入shape[N, C, H, W]完全一致，否则，效果非常非常差，动态模型慎用。

• 操作示例参看onnx2trt_ptq.py

1.2 polygraphy工具:应该是对1.1量化过程的封装

• 操作流程，按照常规方案导出onnx，onnx序列化为tensorrt engine之前打开int8量化模式并采用校正数据集进行校正；

• 优点：1. 相较于1.1，代码量更少，只需完成校正数据的处理代码；

• 缺点：1. 同上所有; 2. 动态尺寸时，校正数据需与–trt-opt-shapes相同；

• 安装polygraphy

  pip install colored polygraphy --extra-index-url https://pypi.ngc.nvidia.com
  

• 量化

polygraphy convert XX.onnx --int8 --data-loader-script loader_data.py --calibration-cache XX.cache -o XX.plan --trt-min-shapes images:[1,3,384,1280] --trt-opt-shapes images:[26,3,384,1280] --trt-max-shapes images:[26,3,384,1280] #量化

• loader_data.py为较正数据集加载过程，自动调用脚本中的load_data()函数：

2. pytorch中执行(推荐)

注：在pytorch中执行导出的onnx将产生一个明确量化的模型，属于显示量化

• 操作流程：安装pytorch_quantization库->加载校正数据->加载模型（在加载模型之前，启用quant_modules.initialize() 以保证原始模型层替换为量化层）->校正->导出onnx;

• 优点：1. 通过导出的onnx能够看到每层量化的过程；2. onnx导出为tensort engine时可以采用trtexec(注：命令行需加–int8，需要fp16和int8混合精度时，再添加–fp16)，比较简单；3. pytorch校正过程可在任意设备中进行；4.相较上述方法，校正数据集使用shape无需与推理shape一致，也能获得较好的结果，动态输入时，推荐采用此种方式。

• 缺点：导出onnx时，显存占用非常大；

• 操作示例参看：pytorch模型进行量化导出yolov5_pytorch_ptq.py

方式3：QAT(追求精度时推荐)

注：在pytorch中执行导出的onnx将产生一个明确量化的模型，属于显式量化

• 操作流程：安装pytorch_quantization库->加载训练数据->加载模型（在加载模型之前，启用quant_modules.initialize() 以保证原始模型层替换为量化层）->训练->导出onnx;

• 优点：1. 模型量化参数重新训练，训练较好时，精度下降较少； 2. 通过导出的onnx能够看到每层量化的过程；2. onnx导出为tensort engine时可以采用trtexec(注：命令行需加–int8，需要fp16和int8混合精度时，再添加–fp16)，比较简单；3.训练过程可在任意设备中进行；

• 缺点：1.导出onnx时，显存占用非常大；2.最终精度取决于训练好坏；3. QAT训练shape需与推理shape一致才能获得好的推理结果；4. 导出onnx时需采用真实的图片输入作为输入设置

• 操作示例参看yolov5_pytorch_qat.py感知训练，参看export_onnx_qat.py

总结

• int8量化对小目标检测影响较大；
• int8量化相比fp16量化推理时间并不会节省一半，需实测；
• 当fp16推理时间满足要求时，请采用fp16量化；