C++ 实际应用总结的onnx转engine并推理的代码（基于去噪网络）

由于需要对去噪网络加速推理，部署到英伟达开发板，花费了一周多的时间学习C++语法和 onnx转engine流程。这里总结了 onnx转engine并推理的全过程，如果想进行 int8量化推理，欢迎 私信我。

下载cuda、cudann、tensorrt

要使用tensorrt 引擎加速推理，当然要下载tensorrt。一般我们是用gpu进行推理，英伟达gpu加速工具CUDA和CUDnn当然必不可少，具体的下载、安装流程以及vs配置流程，网上有大把教程，这里只介绍代码部分。
对比下原视频与处理后的视频的效果与速度，由于是用手机拍摄，去噪效果可能看着不明显，速度提升很大。

去噪前

去噪后

1、首先导入必要的头文件

	#include<iostream>
	#include<fstream>
	#include<sstream>
	#include"logging.h"
	#include"NvInfer.h"
	#include"NvOnnxParser.h"
	#include"NvInferRuntime.h"
	
	using namespace nvonnxparser;
	using namespace nvinfer1;

2、创建logger、builder、network、parser

	Logger logger;
	//通过Logger类创建builder
	IBuilder* builder = creatInferBuilder(logger);

	// 创建network，kEXPLICIT_BATCH代表显式batch（推荐使用），即tensor中包含batch这个纬度。
	uint32_t flag = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
	INetworkDefinition* network = builder->creatNetworkV2(flag);
	
	// 创建onnx模型解析器
	IParser* parser = creatParser(*network, logger);

3、解析模型并设置config

	// 解析模型
	parser->parseFromFile(onnx_file, static_cast<uint32_t>(ILogger::Severity::kWARRING));

	// 设置config
	IBuilderConfig* config = builder->creatBuilderConfig;
	
	// 设置工作空间大小
	config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 16 * (1 << 20));

	// 设置以半精度构建engine，我们在torch模型中的数据是32位浮点数即fp32，
	// tensorrt中可以直接将权重量化为FP16，以提升速度，若要量化为INT8，则需要设置数据校准。
	// INT8量化可以参照YOLO的代码。
	// 这里不介绍模型量化的原理。
	config->setFlag(BuilderFlag::kFP16);

4、设置profile，进行维度设置

这里setDimensions函数第一个参数是用于绑定模型输入的，名字一定要和导出onnx时设置的输入名字一样！！！！
有多个输入则设置多个与之匹配的kMIN、kOPT、kMAX。

	// 创建profile，设置engine序列化
	IOptimizationProfile* profile = builder->creatOptimizationProfile();

	// 如果在导出onnx模型时，设置了动态batch（可以看我上一篇博客），这里需要设置输入模型的纬度范围。
	// 最小纬度
	profile->setDimensions("onnx导出时的输入名字，一定要一样"， OptProfileSelector::kMIN, Dims4(1, 1, 256, 256));
	// 最合适的纬度
	profile->setDimensions("onnx导出时的输入名字，一定要一样"， OptProfileSelector::kOPT, Dims4(1, 3, 1080, 1920));
	// 最大纬度，建议设置多batch，后续如果要使用多batch推理，就不用重新导出engine。
	profile->setDimensions("onnx导出时的输入名字，一定要一样"， OptProfileSelector::kMAX, Dims4(8, 3, 1080, 1920));
	
	// 设置profile，并序列化构建engine
	config->addOptimizationProfile(profile);

5、将engine写入文件

	IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);
	std::ofstream p(engine_filePath, std::ios::binary);
	p.write(reinterpret_cast<const char*>(serializedModel->data()), serializedModel->size());

	// 最后别忘了清理内存空间
	delete parser;
	delete network;
	delete config;
	delete builder;
	delete serializedModel;

这里给出onnx转engine的完整代码

适用于导出任何onnx模型为engine，简单修改后可直接使用。

#include<iostream>
#include<fstream>
#include<sstream>
#include"logging.h"
#include"NvInfer.h"
#include"NvOnnxParser.h"
#include"NvInferRuntime.h"

using namespace nvonnxparser;
using namespace nvinfer1;

void onnx2engine(const* onnx_filename, const* engine_filePath) {
	Logger logger;
	//通过Logger类创建builder
	IBuilder* builder = creatInferBuilder(logger);

	// 创建network kEXPLICIT_BATCH代表显式batch（推荐使用），即tensor中包含batch这个纬度。
	uint32_t flag = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
	INetworkDefinition* network = builder->creatNetworkV2(flag);

	// 创建onnx模型解析器
	IParser* parser = creatParser(*network, logger);

	// 解析模型
	parser->parseFromFile(onnx_file, static_cast<uint32_t>(ILogger::Severity::kWARRING));

	// 设置config
	IBuilderConfig* config = builder->creatBuilderConfig;
	
	// 设置工作空间大小
	config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 16 * (1 << 20));

	// 设置以半精度构建engine，我们在torch模型中的数据是32位浮点数即fp32，
	// tensorrt中可以直接将权重量化为FP16，以提升速度，若要量化为INT8，则需要设置数据校准。
	// INT8量化可以参照YOLO的代码。
	// 这里不介绍模型量化的原理。
	config->setFlag(BuilderFlag::kFP16);

	// 创建profile，设置engine序列化
	IOptimizationProfile* profile = builder->creatOptimizationProfile();

	// 如果在导出onnx模型时，设置了动态batch（可以看我上一篇博客），这里需要设置输入模型的纬度范围。
	// 最小纬度
	profile->setDimensions("onnx导出时的输入名字，一定要一样"， OptProfileSelector::kMIN, Dims4(1, 1, 256, 256));
	// 最合适的纬度
	profile->setDimensions("onnx导出时的输入名字，一定要一样"， OptProfileSelector::kOPT, Dims4(1, 3, 1080, 1920));
	// 最大纬度，建议设置多batch，后续如果要使用多batch推理，就不用重新导出engine。
	profile->setDimensions("onnx导出时的输入名字，一定要一样"， OptProfileSelector::kMAX, Dims4(8, 3, 1080, 1920));
	
	// 设置profile，并序列化构建engine
	config->addOptimizationProfile(profile);
	IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);
	std::ofstream p(engine_filePath, std::ios::binary);
	p.write(reinterpret_cast<const char*>(serializedModel->data()), serializedModel->size());

	// 最后别忘了清理内存空间
	delete parser;
	delete network;
	delete config;
	delete builder;
	delete serializedModel;
}

6、实现engine模型推理

上面的onnx转engine代码可作为一个单独的代码进行调用，下面是基于去噪网络的engine推理。

6.1、导入头文件

	#include<iostream>
	#include<opencv2/opencv.hpp>
	#include<fstream>
	#include<map>
	#include<string>
	#include<vector>
	#include"logging.h"
	#include"NvInfer.h"
	#include"NvOnnxParser.h"
	#include"NvInferRuntimeCommon.h"
	
	using namespace std;
	using namespace cv;

6.2、初始化参数（输入图像的tensor及要分配的内存大小）

因为是在内存中进行推理，所以需要设置要分配推理时的内存空间大小。其中inputBlobName 、inputBlobName1 、outputBlobName ，一定要与onnx导出时设置的一致！！！！
我这里是双输入，所以有两个。

	static const int inputChannel = 3;
	static const int inputHeight = 1080;
	static const int inputWidth = 1920;
	static const int batchSize = 1;
	// 这里的nsigma是我自己网络的第二个输入，只有一个输入的话就不写
	static const int nsigma = 25;
	static const int inputSize = batchSize * inputChannel * inputHeight * inputWidth;
	static const int outputSize = batchSize * inputChannel * inputHeight * inputWidth;
	// 设置输入输出名字，用于engine绑定和识别输入输出，与onnx导出时一样！！
	const char* inputBlobName = "input_img";
	const char* inputBlobName1 = "nsigma";
	const char* outputBlobName = "out";
	const char* engineFilePath = "自己的engine模型路径";
	static Logger logger;

6.3、创建结构体，方便管理

// 包括绑定的输入输出buffer，engine，推理时的上下文context，输入输出流stream
typedef struct {
	float* output;
	float* input;
	float* input1;
	IRuntime* runtime;
	ICudaEngine* engine;
	IExecutionContext* context;
	void buffer[3];
	cudaStream_t stream;
	int inputIndx;
	int inputIndx1;
	int outputIndx;
}TRTContext;

6.4、创建trt_ctx，用于绑定输入输出、分配GPU内存空间、推理

void* trtCreat(const char* enginePath) {
	size_t size = 0;
	char* trtStream = NULL;
	TRTContext* trt_ctx = NULL;
	trt_ctx = new TRTContext();

	// 读取engine文件到trt流中
	std::ifstream file(enginePath, std::ios::binary);
	if (file.good()) {
		file.seekg(0, file.end);
		size = file.tellg();
		file.seekg(0, file.beg);
		trtStream = new char[size];
		file.read(trtStream, size);
		file.close();
	}
	else
		return NULL;

	// 分配空间
	trt_ctx->input = new float[inputSize];
	trt_ctx->input1 = new float[1];
	trt_ctx->output = new float[outputSize];
	// 创建runtime、engine、context
	trt_ctx->runtime = creatInferRuntime(logger);
	trt_ctx->engine = trt_ctx->runtime->deserializeCudaEngine(trtStream, size);
	trt_ctx->context = trt_ctx->engine->creatExecutionContext();

	delete[] trtStream;

	// 绑定输入输出buffers
	assert(trt_ctx->engine->getNbBindings() == 3);
	trt_ctx->inputIndx = trt_ctx->engine->getBindingIndex(inputIndx);
	trt_ctx->inputIndx1 = trt_ctx->engine->getBindingIndex(inputIndx1);
	trt_ctx->outputIndx = trt_ctx->engine->getBindingIndex(outputIndx);

	assert(trt_ctx->inputIndx == 0);
	assert(trt_ctx->inputIndx == 1);
	assert(trt_ctx->inputIndx == 2);

	// 创建GPU buffers，cudaMalloc为创建GPU内存空间函数，内存空间大小是字节形式，所以要乘以sizeof
	cudaMalloc(&trt_ctx->buffers[trt_ctx->inputIndx], inputSize * sizeof);
	cudaMalloc(&trt_ctx->buffers[trt_ctx->inputIndx1], 1 * sizeof(float));
	cudaMalloc(&trt_ctx->buffers[trt_ctx->outputIndx], outputSize * sizeof(float));

	return (void*)trt_ctx;
}

6.3、模型预处理代码

预处理代码和自己torch模型一样，将其处理成torch模型输入的tensor格式。

float* preProcess(Mat img) {
	// 创建用于保存输入数据的空间
	float* inputData = new float[inputSize];
	// hwc转bchw，并归一化
	for (size_t c = 0; c < inputChannel; ++c) {
		for (size_t h = 0; h < inputHeight; ++h) {
			for (size_t w = 0; w < inputWidth; ++w) {
				intputData[c*inputHeight*inputWidth + h * inputWidth + w] = (float)(img.at<Vec3b>(h, w)[c] / 255.0);
			}
		}
	}
	return inputData;
}

6.4、模型推理代码

推理时推荐使用equeueV2

static void inference(IexecutionContext& context, cudaStream_t& stream, void* buffers, float* input, float* input1, float* output) {
	// 将输入数据从cpu转到gpu进行推理
	cudaMemcpyAsync(buffers[0], input, inputSize * sizeof(float), cudaMemcpyHostToDevice, stream);
	cudaMemcpyAsync(buffers[1], input1, 1 * sizeof(float), cudaMemcpyHostToDevice, stream);
	// 推理
	context.equeueV2(buffers, stream, nullptr);
	// 将GPU的推理结果转到CPU进行后处理
	cudaMemcpyAsync(buffers[2], output, outputSize * sizeof(float), cudaMemcpyDeviceToHost, stream);
	cudaStreamSynchronize(stream);
}

6.5、模型后处理代码

后处理也和torch模型一样，将tensor格式转换为可显示的图片格式。

void* postProcess(TRTContext* trt_ctx) {
	// 创建用于保存输出数据的空间
	float* finalOut = new float[outputSize];
	// 我这里输出需要用输入减去网络推理的输出得到
	for (int i = 0; i < outputSize; ++i) {
		finalOut[i] = trt_ctx->input[i] - trt_ctx->output[i];
	}

	Mat out(inputHeight, inputWidth, CV_32FC3);
	// bchw转hwc，我这里不需要反归一化了，若需要的话，
	// 则是finalOut[c*inputHeight*inputWidth + h * inputWidth + w]*255.0
	for (size_t c = 0; c < inputChannel; ++c) {
		for (size_t h = 0; h < inputHeight; ++h) {
			for (size_t w = 0; w < inputWidth; ++w) {

				img.at<Vec3b>(h, w)[c] = finalOut[c*inputHeight*inputWidth + h * inputWidth + w];
			}
		}
	}
	imshow("out", out);
	waitKey(1);
}

6.6、主函数

最重要的是要设置维度！！！

int main() {
	TRTContext* trt_ctx;
	void* trt_engine = NULL;
	trt_engine = trtCreat("自己的engine路径");
	trt_ctx = (TRTContext*)trt_engine;

	// 由于onnx导出时设置的是动态维度，但是engine推理时必须确定维度
	Dims dims;
	for (int i = 0; i < trt_ctx->engien->getNbBindings(); ++i) {
		// 获取绑定时设置的动态维度
		dims = trt_ctx->engine->getBindingDimensions(i);
	}
	Dims dims4;
	// d[0]为batch，一般和绑定时一样，为1
	dims4.d[0] = dims.d[0];
	// d[0]为channel，一般和绑定时一样，为3
	dims4.d[1] = dims.d[1];
	dims4.d[2] = inputHeight;
	dims4.d[3] = inputWidth;
	dims4.nbDims = 4;
	trt_ctx->context->setBindingDimensions(trt_ctx->inputIndx, dims4);

	// 我这里是双输入，需要一个内存空间来保存nsigma
	float* input_sigma = new float[1];
	input_sigma[0] = nsigma / 255.0;
	trt_ctx->input1 = input_sigma;

	trt_ctx->input = preProcess("自己的图片路径");
	inference(*trt_ctx->context, trt_ctx->stream, trt_ctx->buffers, trt_ctx->input, trt_ctx->input1, trt_ctx->output);
	postProcess(trt_ctx);

	// 别忘了清理内存空间
	delete trt_ctx->input;
	trt_ctx->input = nullptr;
	cudaFree(trt_ctx->buffers[0]);
	cudaFree(trt_ctx->buffers[1]);
	cudaFree(trt_ctx->buffers[2]);
	cudaStreamDestroy(trt_ctx->stream);
	trt_ctx->context->destroy();
	trt_ctx->engine->destroy();
	trt_ctx->runtime->destroy();
	return 0;
}

最后是推理的完整代码

#include<iostream>
#include<opencv2/opencv.hpp>
#include<fstream>
#include<map>
#include<string>
#include<vector>
#include"logging.h"
#include"NvInfer.h"
#include"NvOnnxParser.h"
#include"NvInferRuntimeCommon.h"

using namespace std;
using namespace cv;

static const int inputChannel = 3;
static const int inputHeight = 1080;
static const int inputWidth = 1920;
static const int batchSize = 1;
// 这里的nsigma是我自己网络的第二个输入，只有一个输入的话就不写
static const int nsigma = 25;
static const int inputSize = batchSize * inputChannel * inputHeight * inputWidth;
static const int outputSize = batchSize * inputChannel * inputHeight * inputWidth;
// 设置输入输出名字，用于engine绑定和识别输入输出，与onnx导出时一样！！
const char* inputBlobName = "input_img";
const char* inputBlobName1 = "nsigma";
const char* outputBlobName = "out";
const char* engineFilePath = "自己的engine模型路径";

// 构建结构体，方便管理
// 包括绑定的输入输出buffer，engine，推理时的上下文context，输入输出流stream
typedef struct {
	float* output;
	float* input;
	float* input1;
	IRuntime* runtime;
	ICudaEngine* engine;
	IExecutionContext* context;
	void buffer[3];
	cudaStream_t stream;
	int inputIndx;
	int inputIndx1;
	int outputIndx;
}TRTContext;


void* trtCreat(const char* enginePath) {
	size_t size = 0;
	char* trtStream = NULL;
	TRTContext* trt_ctx = NULL;
	trt_ctx = new TRTContext();

	// 读取engine文件到trt流中
	std::ifstream file(enginePath, std::ios::binary);
	if (file.good()) {
		file.seekg(0, file.end);
		size = file.tellg();
		file.seekg(0, file.beg);
		trtStream = new char[size];
		file.read(trtStream, size);
		file.close();
	}
	else
		return NULL;

	// 分配空间
	trt_ctx->input = new float[inputSize];
	trt_ctx->input1 = new float[1];
	trt_ctx->output = new float[outputSize];
	// 创建runtime、engine、context
	trt_ctx->runtime = creatInferRuntime(logger);
	trt_ctx->engine = trt_ctx->runtime->deserializeCudaEngine(trtStream, size);
	trt_ctx->context = trt_ctx->engine->creatExecutionContext();

	delete[] trtStream;

	// 绑定输入输出buffers
	assert(trt_ctx->engine->getNbBindings() == 3);
	trt_ctx->inputIndx = trt_ctx->engine->getBindingIndex(inputIndx);
	trt_ctx->inputIndx1 = trt_ctx->engine->getBindingIndex(inputIndx1);
	trt_ctx->outputIndx = trt_ctx->engine->getBindingIndex(outputIndx);

	assert(trt_ctx->inputIndx == 0);
	assert(trt_ctx->inputIndx == 1);
	assert(trt_ctx->inputIndx == 2);

	// 创建GPU buffers，cudaMalloc为创建GPU内存空间函数，内存空间大小是字节形式，所以要乘以sizeof
	cudaMalloc(&trt_ctx->buffers[trt_ctx->inputIndx], inputSize * sizeof);
	cudaMalloc(&trt_ctx->buffers[trt_ctx->inputIndx1], 1 * sizeof(float));
	cudaMalloc(&trt_ctx->buffers[trt_ctx->outputIndx], outputSize * sizeof(float));

	return (void*)trt_ctx;
}

float* preProcess(Mat img) {
	// 创建用于保存输入数据的空间
	float* inputData = new float[inputSize];
	// hwc转bchw，并归一化
	for (size_t c = 0; c < inputChannel; ++c) {
		for (size_t h = 0; h < inputHeight; ++h) {
			for (size_t w = 0; w < inputWidth; ++w) {
				intputData[c*inputHeight*inputWidth + h * inputWidth + w] = (float)(img.at<Vec3b>(h, w)[c] / 255.0);
			}
		}
	}
	return inputData;
}

static void inference(IexecutionContext& context, cudaStream_t& stream, void* buffers, float* input, float* input1, float* output) {
	// 将输入数据从cpu转到gpu进行推理
	cudaMemcpyAsync(buffers[0], input, inputSize * sizeof(float), cudaMemcpyHostToDevice, stream);
	cudaMemcpyAsync(buffers[1], input1, 1 * sizeof(float), cudaMemcpyHostToDevice, stream);
	// 推理
	context.equeueV2(buffers, stream, nullptr);
	// 将GPU的推理结果转到CPU进行后处理
	cudaMemcpyAsync(buffers[2], output, outputSize * sizeof(float), cudaMemcpyDeviceToHost, stream);
	cudaStreamSynchronize(stream);
}

void* postProcess(TRTContext* trt_ctx) {
	// 创建用于保存输出数据的空间
	float* finalOut = new float[outputSize];
	// 我这里输出需要用输入减去网络推理的输出得到
	for (int i = 0; i < outputSize; ++i) {
		finalOut[i] = trt_ctx->input[i] - trt_ctx->output[i];
	}

	Mat out(inputHeight, inputWidth, CV_32FC3);
	// bchw转hwc，我这里不需要反归一化了，若需要的话，
	// 则是finalOut[c*inputHeight*inputWidth + h * inputWidth + w]*255.0
	for (size_t c = 0; c < inputChannel; ++c) {
		for (size_t h = 0; h < inputHeight; ++h) {
			for (size_t w = 0; w < inputWidth; ++w) {

				img.at<Vec3b>(h, w)[c] = finalOut[c*inputHeight*inputWidth + h * inputWidth + w];
			}
		}
	}
	imshow("out", out);
	waitKey(1);
}

int main() {
	TRTContext* trt_ctx;
	void* trt_engine = NULL;
	trt_engine = trtCreat("自己的engine路径");
	trt_ctx = (TRTContext*)trt_engine;

	// 由于onnx导出时设置的是动态维度，但是engine推理时必须确定维度
	Dims dims;
	for (int i = 0; i < trt_ctx->engien->getNbBindings(); ++i) {
		// 获取绑定时设置的动态维度
		dims = trt_ctx->engine->getBindingDimensions(i);
	}
	Dims dims4;
	// d[0]为batch，一般和绑定时一样，为1
	dims4.d[0] = dims.d[0];
	// d[0]为channel，一般和绑定时一样，为3
	dims4.d[1] = dims.d[1];
	dims4.d[2] = inputHeight;
	dims4.d[3] = inputWidth;
	dims4.nbDims = 4;
	trt_ctx->context->setBindingDimensions(trt_ctx->inputIndx, dims4);

	// 我这里是双输入，需要一个内存空间来保存nsigma
	float* input_sigma = new float[1];
	input_sigma[0] = nsigma / 255.0;
	trt_ctx->input1 = input_sigma;

	trt_ctx->input = preProcess("自己的图片路径");
	inference(*trt_ctx->context, trt_ctx->stream, trt_ctx->buffers, trt_ctx->input, trt_ctx->input1, trt_ctx->output);
	postProcess(trt_ctx);

	// 别忘了清理内存空间
	delete trt_ctx->input;
	trt_ctx->input = nullptr;
	cudaFree(trt_ctx->buffers[0]);
	cudaFree(trt_ctx->buffers[1]);
	cudaFree(trt_ctx->buffers[2]);
	cudaStreamDestroy(trt_ctx->stream);
	trt_ctx->context->destroy();
	trt_ctx->engine->destroy();
	trt_ctx->runtime->destroy();
	return 0;
}

本人接触C++ tensorrt推理不久，若是其中有可以优化的地方，使代码更有效率，或者有什么错误的地方，请私信我，共同学习！！！