TensorRT模型量化实践

本文主要是介绍TensorRT模型量化实践，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

文章目录

- 量化基本概念
- 量化的方法
- - 方式1：trtexec（PTQ的一种）
  - 方式2：PTQ
  - - 2.1 python onnx转trt
    - 2.2 polygraphy工具:应该是对2.1量化过程的封装
  - 方式3：QAT(追求精度时推荐)
- 使用TensorRT量化实践（C++版）
- 使用TensorRT量化（python版）
- 参考文献

量化基本概念

后训练量化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.assign(path).append("\\*").c_str(), &FileInfo);while (_findnext(Handle, &FileInfo) == 0) {if (strcmp(FileInfo.name, ".") != 0 && strcmp(FileInfo.name, "..") != 0) {files.push_back(FileInfo.name);}}
}Calibrator::Calibrator(int batchsize, int input_w, int input_h, std::string img_dir, const char* calib_table_name, bool read_cache){BATCHSIZE = batchsize;WIDTH = input_w;HEIGHT = input_h;INDEX = 0;IMAGEDIR = img_dir;CALIBRATORTABLE = calib_table_name;READCACHE = read_cache;INPUTSIZE = BATCHSIZE * 3 * WIDTH * HEIGHT;cudaMalloc(&DEVICEINPUT, INPUTSIZE * sizeof(float));getFiles(IMAGEDIR, IMAGEFILES);
}Calibrator::~Calibrator() {cudaFree(DEVICEINPUT);
}int Calibrator::getBatchSize() const noexcept {return BATCHSIZE;
}bool Calibrator::getBatch(void* bindings[], const char* names[], int nbBindings) noexcept {if (INDEX + BATCHSIZE > (int)IMAGEFILES.size()) return false;std::vector<cv::Mat> input_imgs;for (int i = INDEX; i < INDEX + BATCHSIZE; i++) {cv::Mat temp = cv::imread(IMAGEDIR + IMAGEFILES[i]);if (temp.empty()) {std::cerr << "Image cannot open!" << std::endl;return false;}cv::Mat pr_img = preprocess_img(temp, WIDTH, HEIGHT);input_imgs.push_back(pr_img);}INDEX += BATCHSIZE;cv::Mat blob = cv::dnn::blobFromImages(input_imgs, 1.0 / 255.0, cv::Size(WIDTH, HEIGHT), cv::Scalar(0, 0, 0), true, false);cudaMemcpy(DEVICEINPUT, blob.ptr<float>(0), INPUTSIZE * sizeof(float), cudaMemcpyHostToDevice);bindings[0] = DEVICEINPUT;return true;
}const void* Calibrator::readCalibrationCache(size_t& length) noexcept {std::cout << "reading calib cache: " << CALIBRATORTABLE << std::endl;CALIBRATORCACHE.clear();std::ifstream input(CALIBRATORTABLE, std::ios::binary);input >> std::noskipws;if (READCACHE && input.good()) {std::copy(std::istream_iterator<char>(input), std::istream_iterator<char>(), std::back_inserter(CALIBRATORCACHE));}length = CALIBRATORCACHE.size();return length ? CALIBRATORCACHE.data() : nullptr;
}void Calibrator::writeCalibrationCache(const void* cache, size_t length) noexcept {std::cout << "writing calib cache: " << CALIBRATORTABLE << std::endl;std::ofstream output(CALIBRATORTABLE, std::ios::binary);output.write(reinterpret_cast<const char*>(cache), length);
}

最后，通过以下代码将onnx量化转换为engine文件。

#include <iostream>
#include <fstream>
#include "calibrator.h"
#include "NvInfer.h"
#include "NvOnnxParser.h"// 实例化记录器界面，捕获所有警告性信息，但忽略信息性消息
class Logger : public nvinfer1::ILogger {void log(Severity severity, const char* msg) noexcept override {if (severity <= Severity::kWARNING) {std::cout << msg << std::endl;}}
}logger;void ONNX2TensorRT(const char* ONNX_file, std::string& Engine_file, bool& FP16, bool& INT8, std::string& image_dir, const char*& calib_table) {std::cout << "Load ONNX file form: " << ONNX_file << "\nStart export..." << std::endl;// 1.创建构建器的实例nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);// 2.创建网络定义uint32_t flag = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);nvinfer1::INetworkDefinition* network = builder->createNetworkV2(flag);// 3.创建一个 ONNX 解析器来填充网络nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger);// 4.读取模型文件并处理任何错误parser->parseFromFile(ONNX_file, static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING));for (int32_t i = 0; i < parser->getNbErrors(); ++i)std::cout << parser->getError(i)->desc() << std::endl;// 5.创建构建配置，指定TensorRT如何优化模型nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();// 如果是动态模型，则需要设置大小/*auto profile = builder->createOptimizationProfile();auto input_tensor = network->getInput(0);auto input_dims = input_tensor->getDimensions();// 配置最小允许batchinput_dims.d[0] = 1;profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, input_dims);profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, input_dims);// 配置最大允许batch// if networkDims.d[i] != -1, then minDims.d[i] == optDims.d[i] == maxDims.d[i] == networkDims.d[i]input_dims.d[0] = maxBatchSize;profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);config->addOptimizationProfile(profile);*/// 6.设置属性来控制 TensorRT 如何优化网络// 设置内存池的空间config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 16 * (1 << 20));if (FP16) {// 判断硬件是否支持FP16if (!builder->platformHasFastFp16()) {std::cout << "不支持FP16量化！" << std::endl;system("pause");return;}config->setFlag(nvinfer1::BuilderFlag::kFP16);}else if (INT8) {if (!builder->platformHasFastInt8()) {std::cout << "不支持INT8量化！" << std::endl;system("pause");return;}config->setFlag(nvinfer1::BuilderFlag::kINT8);nvinfer1::IInt8EntropyCalibrator2* calibrator = new Calibrator(1, 640, 640, image_dir, calib_table);config->setInt8Calibrator(calibrator);}// 7.指定配置后，构建引擎nvinfer1::IHostMemory* serializeModel = builder->buildSerializedNetwork(*network, *config);// 8.保存TensorRT模型std::ofstream engine(Engine_file, std::ios::binary);engine.write(reinterpret_cast<const char*>(serializeModel->data()), serializeModel->size());// 9.序列化引擎包含权重的必要副本，因此不再需要解析器、网络定义、构建器配置和构建器，可以安全地删除delete parser;delete network;delete config;delete builder;// 10.将引擎保存到磁盘后 ，并且可以删除它被序列化到的缓冲区delete serializeModel;std::cout << "Export success, Save as: " << Engine_file << std::endl;
}int main(int argc, char** argv) {// ONNX 文件路径const char* ONNX_file = "../weights/yolov8s.onnx";// ENGINE 文件保存路径std::string Engine_file = "../weights/yolov8s.engine";// 当量化为INT8时，图片路径std::string image_dir = "../images/";// 当量化为INT8时，校准表路径（存在读取，不存在创建）const char* calib_table = "../weights/calibrator.table";// 选择量化方式，若两个都为false，使用FP32生成 ENGINE文件bool FP16 = false;bool INT8 = true;std::ifstream file(ONNX_file, std::ios::binary);if (!file.good()) {std::cout << "Load ONNX file failed！" << std::endl;}ONNX2TensorRT(ONNX_file, Engine_file, FP16, INT8, image_dir, calib_table);return 0;
}

使用TensorRT量化（python版）

流程C++版本的一样，这个没进行测试，以下版本是别人量化yolov5的代码，感兴趣的朋友可以尝试一下。

import tensorrt as trt
import os
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import cv2def get_crop_bbox(img, crop_size):"""Randomly get a crop bounding box."""margin_h = max(img.shape[0] - crop_size[0], 0)margin_w = max(img.shape[1] - crop_size[1], 0)offset_h = np.random.randint(0, margin_h + 1)offset_w = np.random.randint(0, margin_w + 1)crop_y1, crop_y2 = offset_h, offset_h + crop_size[0]crop_x1, crop_x2 = offset_w, offset_w + crop_size[1]return crop_x1, crop_y1, crop_x2, crop_y2def crop(img, crop_bbox):"""Crop from ``img``""" crop_x1, crop_y1, crop_x2, crop_y2 = crop_bboximg = img[crop_y1:crop_y2, crop_x1:crop_x2, ...]return imgclass yolov5EntropyCalibrator(trt.IInt8EntropyCalibrator2):def __init__(self, imgpath, batch_size, channel, inputsize=[384, 1280]):trt.IInt8EntropyCalibrator2.__init__(self)self.cache_file = 'yolov5.cache'self.batch_size = batch_sizeself.Channel = channelself.height = inputsize[0]self.width = inputsize[1]self.imgs = [os.path.join(imgpath, file) for file in os.listdir(imgpath) if file.endswith('jpg')]np.random.shuffle(self.imgs)self.imgs = self.imgs[:2000]self.batch_idx = 0self.max_batch_idx = len(self.imgs) // self.batch_sizeself.calibration_data = np.zeros((self.batch_size, 3, self.height, self.width), dtype=np.float32)# self.data_size = trt.volume([self.batch_size, self.Channel, self.height, self.width]) * trt.float32.itemsizeself.data_size = self.calibration_data.nbytesself.device_input = cuda.mem_alloc(self.data_size)# self.device_input = cuda.mem_alloc(self.calibration_data.nbytes)def free(self):self.device_input.free()def get_batch_size(self):return self.batch_sizedef get_batch(self, names, p_str=None):try:batch_imgs = self.next_batch()if batch_imgs.size == 0 or batch_imgs.size != self.batch_size * self.Channel * self.height * self.width:return Nonecuda.memcpy_htod(self.device_input, batch_imgs)return [self.device_input]except:print('wrong')return Nonedef next_batch(self):if self.batch_idx < self.max_batch_idx:batch_files = self.imgs[self.batch_idx * self.batch_size: \(self.batch_idx + 1) * self.batch_size]batch_imgs = np.zeros((self.batch_size, self.Channel, self.height, self.width),dtype=np.float32)for i, f in enumerate(batch_files):img = cv2.imread(f)  # BGRcrop_size = [self.height, self.width]crop_bbox = get_crop_bbox(img, crop_size)# crop the imageimg = crop(img, crop_bbox)img = img.transpose((2, 0, 1))[::-1, :, :]  # BHWC to BCHW ,BGR to RGBimg = np.ascontiguousarray(img)img = img.astype(np.float32) / 255.assert (img.nbytes == self.data_size / self.batch_size), 'not valid img!' + fbatch_imgs[i] = imgself.batch_idx += 1print("batch:[{}/{}]".format(self.batch_idx, self.max_batch_idx))return np.ascontiguousarray(batch_imgs)else:return np.array([])def read_calibration_cache(self):# If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.if os.path.exists(self.cache_file):with open(self.cache_file, "rb") as f:return f.read()def write_calibration_cache(self, cache):with open(self.cache_file, "wb") as f:f.write(cache)f.flush()# os.fsync(f)def get_engine(onnx_file_path, engine_file_path, cali_img, mode='FP32', workspace_size=4096):"""Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""TRT_LOGGER = trt.Logger(trt.Logger.WARNING)def build_engine():assert mode.lower() in ['fp32', 'fp16', 'int8'], "mode should be in ['fp32', 'fp16', 'int8']"explicit_batch_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)with trt.Builder(TRT_LOGGER) as builder, builder.create_network(explicit_batch_flag) as network, builder.create_builder_config() as config, trt.OnnxParser(network, TRT_LOGGER) as parser:with open(onnx_file_path, "rb") as model:print("Beginning ONNX file parsing")if not parser.parse(model.read()):print("ERROR: Failed to parse the ONNX file.")for error in range(parser.num_errors):print(parser.get_error(error))return Noneconfig.max_workspace_size = workspace_size * (1024 * 1024)  # workspace_sizeMiB# 构建精度if mode.lower() == 'fp16':config.flags |= 1 << int(trt.BuilderFlag.FP16)if mode.lower() == 'int8':print('trt.DataType.INT8')config.flags |= 1 << int(trt.BuilderFlag.INT8)config.flags |= 1 << int(trt.BuilderFlag.FP16)calibrator = yolov5EntropyCalibrator(cali_img, 26, 3, [384, 1280])# config.set_quantization_flag(trt.QuantizationFlag.CALIBRATE_BEFORE_FUSION)config.int8_calibrator = calibrator# if True:#     config.profiling_verbosity = trt.ProfilingVerbosity.DETAILEDprofile = builder.create_optimization_profile()profile.set_shape(network.get_input(0).name, min=(1, 3, 384, 1280), opt=(12, 3, 384, 1280), max=(26, 3, 384, 1280))config.add_optimization_profile(profile)# config.set_calibration_profile(profile)print("Completed parsing of ONNX file")print("Building an engine from file {}; this may take a while...".format(onnx_file_path))# plan = builder.build_serialized_network(network, config)# engine = runtime.deserialize_cuda_engine(plan)engine = builder.build_engine(network,config)print("Completed creating Engine")with open(engine_file_path, "wb") as f:# f.write(plan)f.write(engine.serialize())return engineif os.path.exists(engine_file_path):# If a serialized engine exists, use it instead of building an engine.print("Reading engine from file {}".format(engine_file_path))with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:return runtime.deserialize_cuda_engine(f.read())else:return build_engine()def main(onnx_file_path, engine_file_path, cali_img_path, mode='FP32'):"""Create a TensorRT engine for ONNX-based YOLOv3-608 and run inference."""# Try to load a previously generated YOLOv3-608 network graph in ONNX format:get_engine(onnx_file_path, engine_file_path, cali_img_path, mode)if __name__ == "__main__":onnx_file_path = '/home/models/boatdetect_yolov5/last_nms_dynamic.onnx'engine_file_path = "/home/models/boatdetect_yolov5/last_nms_dynamic_onnx2trtptq.plan"cali_img_path = '/home/data/frontview/test'main(onnx_file_path, engine_file_path, cali_img_path, mode='int8')

参考文献

tensorrt官方int8量化方法汇总
深度学习模型量化基础
模型量化5：onnx模型的静态量化和动态量化
有用的 模型量化！ONNX转TensorRT(FP32, FP16, INT8)
TensorRT-Int8量化详解
TensorRT中的INT 8 优化
TensorRT——INT8推理
TensorRT模型，INT8量化Python实践教程

这篇关于TensorRT模型量化实践的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！