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}) # 链接库

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;
}