本文为C++实现,Python实现链接模型量化
支持三种不同精度的量化
模型单精度量化 (FP32)
模型半精度量化 (FP16)
模型Int8量化 (INT8)
经测试yolov5,yolov6,yolov7,yolov8转化成功
yolov5: https://github.com/ultralytics/yolov5
yolov6: https://github.com/meituan/YOLOv6
yolov7: https://github.com/WongKinYiu/yolov7
yolov8: https://github.com/ultralytics/ultralytics
注意:若使用INT8量化,需要额外的文件,这里是calibrator.cpp & calibrator.h
文章目录
先决条件
如果还没有安装CUDA和CUDNN,请参考这篇文章:
文章链接:Win10配置CUDA和cudnn
我的CUDA版本:11.3
我的CUDNN版本:8.4.1
下载TensorRT
建议下载GA版本(正式版本),链接:https://developer.nvidia.com/tensorrt
我下载的8.4.1.5版本,下载完成后解压
解压后可以看到这三个文件夹
Debug模式
1. 创建新项目
2. 复制TensorRT文件至CUDA
将lib中的dll文件复制到 …\NVIDIA GPU Computing Toolkit\CUDA\v11.3\bin(这是我的路径)
3. 创建onnx2TensorRT函数
当前是Debug模式,将代码copy进去,这时出现错误(红色波浪线)不用管。
main.cpp
#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();
// 6.设置属性来控制 TensorRT 如何优化网络
// 设置内存池的空间
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 16 * (1 << 20));
if (FP16) {
// 判断硬件是否支持FP16
if (!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;
}
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);
}
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);
};
4. 设置属性
右键属性
4.1. 设置属性 — 包含目录
找到VC++目录中的包含目录,将下面的路径写入,保存退出
注意:TensorRT路径需要替换成你的路径,CUDA路径复制即可。
$(CUDA_PATH)\include
D:\software\opencv\build\include
D:\software\TensorRT-8.4.1.5_CUDA11.6_Cudnn8.4.1\include
D:\software\TensorRT-8.4.1.5_CUDA11.6_Cudnn8.4.1\samples\common
4.2. 设置属性 — 库目录
找到库目录,将下面的路径写入,保存退出
注意:TensorRT路径需要替换成你的路径,CUDA路径复制即可。
$(CUDA_PATH)\lib
$(CUDA_PATH)\lib\x64
D:\software\opencv\build\x64\vc15\lib
D:\software\TensorRT-8.4.1.5_CUDA11.6_Cudnn8.4.1\lib
4.3. 设置属性 — 链接器
找到链接器下的输入,在附加依赖项中加入
cudnn.lib
cublas.lib
cudart.lib
nvinfer.lib
nvparsers.lib
nvonnxparser.lib
nvinfer_plugin.lib
opencv_world460d.lib
5. 创建模型文件夹
新建weights文件夹,将你的onnx文件放在里面
6. 运行
点击运行即可(可能有点慢)
7. 问题 — 缺少 zlibwapi.dll 文件
若出现缺少zlibwapi.dll,需要下载,
链接:https://pan.baidu.com/s/12sVdiDH-NOOZNI9QqJoZuA?pwd=a0n0
提取码:a0n0
下载后解压,将zlibwapi.dll 放入 …\NVIDIA GPU Computing Toolkit\CUDA\v11.3\bin(我的路径)
重新运行程序
8. 生成engine文件
Release模式
点击,更换为Release
重新配置属性
重新配置步骤4,添加库目录,包含目录,附加依赖项,即可。