onnx转engine_onnx 转为 engine

CMakeLists.txt

cmake_minimum_required(VERSION 3.0)		# CMake最低版本要求，低于2.6的构建过程会被终止set(CMAKE_CXX_STANDARD 14)		#opencv4以上加此句
set(CMAKE_CXX_STANDARD  14)
project(onnx2engine)				# 定义工程名称
find_package(OpenCV     REQUIRED)
# find_package(PCL 	    REQUIRED)
find_package(CUDA 	    REQUIRED)
find_package(TensorRT 	REQUIRED)
include_directories(
                    # ${PCL_INCLUDE_DIRS} 
                    ${CUDA_INCLUDE_DIRS} 
                    ${TensorRT_INCLUDE_DIRS})
add_executable(onnx2engine src/main.cpp)	# 生成一个叫 cmake_test的可执行文件，其中可执行源文件路径为 src/main.cpp
target_link_libraries( onnx2engine 
                        ${OpenCV_LIBS} 
                        # ${PCL_LIBRARIES} 
                        ${CUDA_LIBRARIES} 
                        ${TensorRT_LIBS}) # 链接库
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

main.cpp

  
#include <fstream>
#include <iostream>
#include <sstream>
 
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "NvInferRuntime.h"
 
using namespace nvinfer1;
using namespace nvonnxparser;
 
// 全局创建 ILogger 类型的对象
class Logger : public ILogger
{
	virtual void log(Severity severity, const char* msg) noexcept override
	{
		// suppress info-level messages
		if (severity != Severity::kINFO)
			std::cout << msg << std::endl;
	}
} gLogger;
 
int onnx2engine(std::string onnx_filename, std::string engine_filePath, bool fp16, int workspaceSize,
	bool dynamic, int batch, int channel, int height, int width) {
	// 创建builder
	// 创建日志：TRTLogger logger
	IBuilder* builder = createInferBuilder(gLogger);
 
	// 创建network
	// 创建一个 network，要求网络结构里，没有隐藏的批量处理维度 
	// 创建网络：createNetworkV2(1)： 显性batch（NCHW）；createNetworkV2(0) ：非显性batch（CHW），新版本一般考虑使用createNetworkV2(1)
	// 创建network，kEXPLICIT_BATCH代表显式batch（推荐使用），即tensor中包含batch这个纬度。
	nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
 
	// 创建onnx模型解析器
	auto parser = nvonnxparser::createParser(*network, gLogger);
 
	// 解析模型
	parser->parseFromFile(onnx_filename.c_str(), 2);
	for (int i = 0; i < parser->getNbErrors(); ++i)
	{
		std::cout << parser->getError(i)->desc() << std::endl;
	}
	printf("tensorRT load onnx model sucessful! \n"); 					// 解析模型成功，第一个断点测试
 
	// 添加配置参数, 告诉TensorRT如何优化网络
	IBuilderConfig* config = builder->createBuilderConfig();
	config->setMaxWorkspaceSize(1024*1024*workspaceSize);  			// 设置最大工作空间
	config->setFlag(BuilderFlag::kGPU_FALLBACK);  					// 启用GPU回退模式，作用？
	//	config->setFlag(BuilderFlag::kSTRICT_TYPES);  				//强制执行xx位的精度计算


	// 设置以半精度构建engine，我们在torch模型中的数据是32位浮点数即fp32，
	// tensorrt中可以直接将权重量化为FP16，以提升速度，若要量化为INT8，则需要设置数据校准。
	// INT8量化可以参照YOLO的代码
	// 这里不介绍模型量化的原理
	if (fp16) {
		config->setFlag(nvinfer1::BuilderFlag::kFP16); 						// 设置精度计算
	}
 
	if(dynamic)
	{
		// 创建profile，设置engine序列化
		IOptimizationProfile* profile = builder->createOptimizationProfile(); //创建优化配置文件
		// 配置动态batch
		// 这里setDimensions函数第一个参数是用于绑定模型输入的，名字一定要和导出onnx时设置的输入名字一样！！！！
		// 有多个输入则设置多个与之匹配的kMIN、kOPT、kMAX
		// onnx导出时的输入名字，一定要一样
		// 最小纬度
		profile->setDimensions("images", nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, channel, height, width));  		
		// 最合适的纬度
		profile->setDimensions("images", nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(batch, channel, height, width)); 	
		// 最大纬度，建议设置多batch，后续如果要使用多batch推理，就不用重新导出engine
		profile->setDimensions("images", nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(batch, channel, height, width));
		// 设置profile，并序列化构建engine	
		config->addOptimizationProfile(profile);
	}

 
	// 使用builder对象构建engine
	ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); //创建engine  第二个断点测试

	// 创建执行上下文
	IExecutionContext* context = engine->createExecutionContext();
	if(!context){
		std::cout<<"context create failed"<<std::endl;
		return false;
	}

	std::cout << "try to save engine file now" << std::endl;

	std::ofstream outf(engine_filePath, std::ios::binary);
	if (!outf) {
		std::cerr << "could not open plan output file" << std::endl;
		return 0;
	}
 
	// 序列化
	IHostMemory* modelStream = engine->serialize(); 								// 第三个断点测试
	outf.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size()); // 写入
 
 	// 销毁
	builder->destroy();
	config->destroy();
	context->destroy();
	modelStream->destroy();
	engine->destroy();
	network->destroy();
	parser->destroy();

	outf.close();
	std::cout << "convert onnx model to TensorRT engine model successfully!" << std::endl; // 转换成功，第四个断点测试
 
	return 0;
}
 
int main(int argc, char** argv) {
	bool fp16 = false;
	int workspaceSize = 64;

	bool dynamic = false;
	int batch =  1;
	int channel = 3;
	int height = 640;
	int width = 640;

	onnx2engine("../data/best.onnx", "../data/best.engine", 
		fp16, workspaceSize, dynamic,
		batch, channel, height, width);
	return 0;
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133