TensorRT模型量化实践

量化基本概念

后训练量化Post Training Quantization (PTQ)

量化过程仅仅通过离线推理一些sample数据对权重和激活值进行量化，无需要进行训练微调。

量化感知训练Quantization Aware Training (QAT)

在量化的过程中，对网络进行训练，从而让网络参数能更好地适应量化带来的信息损失。这种方式更加灵活，因此准确性普遍比后训练量化要高。缺点是操作起来不太方便。大多数情况下比训练后量化精度更高，部分场景不一定比部分/混合精度量化好很多。

量化的方法

方式1：trtexec（PTQ的一种）

（1）int8量化

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

如果使用int8量化；量化需要设置calib文件夹；

trtexec 
--onnx=model.onnx 
--minShapes=input:1x1x224x224  
--optShapes=input:2x1x224x224 
--maxShapes=input:10x1x224x224 
--workspace=4096 
--int8 
--best 
--calib=D:\images 
--saveEngine=model.engine 
--buildOnly  

精度损失很大，不建议直接采用。
trtexec 有提供 --calib=接口进行校正，但需要对中间特征进行cache文件保存，比较麻烦，官方文档也是采用上述方式进行int8量化；与fp16的模型在测试集上测试指标，可以看到精度下降非常严重；
（2）int8 fp16混合量化
trtexec --onnx=XX.onnx --saveEngine=model.plan --int8 --fp16 --workspace=4096
测试集上统计指标：相比纯int8量化，效果要好，但是相比fp16，精度下降依然非常严重

方式2：PTQ

engine序列化时执行

2.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

2.2 polygraphy工具:应该是对2.1量化过程的封装

操作流程：
按照常规方案导出onnx，onnx序列化为tensorrt engine之前打开int8量化模式并采用校正数据集进行校正；
优点： 1. 相较于1.1，代码量更少，只需完成校正数据的处理代码；
缺点： 1. 同上所有; 2. 动态尺寸时，校正数据需与–trt-opt-shapes相同；3.内部默认最多校正20个epoch；

安装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.pl

方式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

使用TensorRT量化实践（C++版）

该方式则是利用TensorRT的API将onnx转换engine文件的过程中进行量化，其中需要校准数据（准备一个存放几百张图像的文件夹即可）。为了读取校正图像，需要写一个Int8校正类，如下所示：
calibrator.h

#pragma once
#include <NvInfer.h>
#include<vector>
#include <opencv2/opencv.hpp>
class Calibrator : public nvinfer1::IInt8EntropyCalibrator2 {
   
   
public:
	Calibrator(int batchsize, int input_w, int input_h, std::string img_dir, const char* calib_table_name, bool read_cache = true);

	virtual ~Calibrator();
	int getBatchSize() const noexcept override;
	bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override;
	const void* readCalibrationCache(size_t& length) noexcept override;
	void writeCalibrationCache(const void* cache, size_t length) noexcept override;

private:
	int BATCHSIZE;
	int WIDTH;
	int HEIGHT;
	int INDEX;
	std::string IMAGEDIR;
	std::vector<std::string> IMAGEFILES;
	size_t INPUTSIZE;
	std::string CALIBRATORTABLE;
	bool READCACHE;
	void* DEVICEINPUT;
	std::vector<char> CALIBRATORCACHE;

	cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h);
	void getFiles(std::string path, std::vector<std::string>& files);
};

calibrator.cpp

#include <fstream>
#include <io.h>
#include "calibrator.h"

cv::Mat Calibrator::preprocess_img(cv::Mat& img, int input_w, int input_h) {
   
   
    int w, h, x, y;
    float r_w = input_w / (img.cols * 1.0);
    float r_h = input_h / (img.rows * 1.0);
    if (r_h > r_w) {
   
   
        w = input_w;
        h = r_w * img.rows;
        x = 0;
        y = (input_h - h) / 2;
    }
    else {
   
   
        w = r_h * img.cols;
        h = input_h;
        x = (input_w - w) / 2;
        y = 0;
    }
    cv::Mat re(h, w, CV_8UC3);
    cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR);
    cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128));
    re.copyTo(out(cv::Rect(x, y, re.cols, re.rows)));
    return out;
}

void Calibrator::getFiles(std::string path, std::vector<std::string>& files){
   
   
    intptr_t Handle;
    struct _finddata_t FileInfo;
    std::string p;
    Handle = _findfirst(p