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

C++ onnx转engine并推理全过程解析(基于去噪网络),设置动态维度、多输入cuda流推理

  • 下载cuda、cudann、tensorrt
  • 1、首先导入必要的头文件
  • 2、创建logger、builder、network、parser
  • 3、解析模型并设置config
  • 4、设置profile,进行维度设置
  • 5、将engine写入文件
  • 这里给出onnx转engine的完整代码
  • 6、实现engine模型推理
    • 6.1、导入头文件
    • 6.2、初始化参数(输入图像的tensor及要分配的内存大小)
    • 6.3、创建结构体,方便管理
    • 6.4、创建trt_ctx,用于绑定输入输出、分配GPU内存空间、推理
    • 6.3、模型预处理代码
    • 6.4、模型推理代码
    • 6.5、模型后处理代码
    • 6.6、主函数
  • 最后是推理的完整代码

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

下载cuda、cudann、tensorrt

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

去噪前

去噪后

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

	#include
	#include
	#include
	#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(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
	INetworkDefinition* network = builder->creatNetworkV2(flag);
	
	// 创建onnx模型解析器
	IParser* parser = creatParser(*network, logger);

3、解析模型并设置config

	// 解析模型
	parser->parseFromFile(onnx_file, static_cast(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(serializedModel->data()), serializedModel->size());

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

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

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

#include
#include
#include
#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(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
	INetworkDefinition* network = builder->creatNetworkV2(flag);

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

	// 解析模型
	parser->parseFromFile(onnx_file, static_cast(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(serializedModel->data()), serializedModel->size());

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

6、实现engine模型推理

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

6.1、导入头文件

	#include
	#include
	#include
	#include
	#include
	#include
	#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(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(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
#include
#include
#include
#include
#include
#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(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(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推理不久,若是其中有可以优化的地方,使代码更有效率,或者有什么错误的地方,请私信我,共同学习!!!

你可能感兴趣的:(c++,开发语言,网络,计算机视觉,神经网络)