onnxToTRTModel (engine) 适用yolov5模型

Posted 2023-03-02 爱吃油淋鸡的莫何

本文件夹有用的文件包括：

1 CMakeLists.txt : 其中涉及 cuda 和 TensorRt 和 OpenCV 相关库的路径

​ 也包含mkdir build, cd build， cmake …，make操作之后生成的target名称
​ add_executable(yolo $PROJECT_SOURCE_DIR/onnx2trt.cpp)中的cpp文件名要匹配

2 onnx2trt.cpp ： onnx 转引擎的主要源文件

3 looging.h ： onnx2trt会调用的库（应该是日志相关的一个库）

4 best.onnx ：要转化的 onnx 文件

5 ReadMe ：此文件

其他注意事项：

onnx2trt.cpp 源码的main()函数中 onnxToTRTModel()调用的参数onnx和engine文件必须是
绝对路径(不知道为什么，总之如果相对路径的话就会报错)

程序如何运行：

在当前目录下依次执行如下操作

mkdir build
cd build
cmake ..
make
./yolo

最后的./yolo执行之后会生成engine文件

代码链接及提取码：

链接：https://pan.baidu.com/s/1viSZ_N7YWHZlvy-LKiYOIw 
提取码：onnx

onnx2trt.cpp：

//onnx2trt.cpp
// onnx文件转为engine文件
#include <fstream>
#include <iostream>
#include <sstream>
#include <numeric>
// #include <chrono>
#include <vector>
#include <opencv2/opencv.hpp>
// #include <dirent.h>
#include "NvInfer.h"

#include "NvOnnxParser.h"
// #include "NvInferRuntime.h"
#include "logging.h"
#include "cuda_runtime_api.h"
using namespace nvinfer1;
using namespace std;
// #define USE_FP16

// 全局创建 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;


static Logger gLogger;

const char* INPUT_BLOB_NAME = "images";
const char* OUTPUT_BLOB_NAME = "419";
const int batchSize = 1;
const char* onnx_path = "/home/mec/hlj/onnx2trt/best.onnx";
const char* TrtSaveFileName = "/home/mec/hlj/onnx2trt/best.engine";



void saveToTrtModel(const char * TrtSaveFileName,IHostMemory*trtModelStream)
	
		std::ofstream out(TrtSaveFileName, std::ios::binary);
		if (!out.is_open())
		
		std::cout << "打开文件失败!" <<std:: endl;
		
		out.write(reinterpret_cast<const char*>(trtModelStream->data()), trtModelStream->size());
		out.close();
	


void onnxToTRTModel(const std::string& modelFile,unsigned int maxBatchSize,IHostMemory*& trtModelStream,const char * TrtSaveFileName) 

    int verbosity = (int) nvinfer1::ILogger::Severity::kWARNING;
 
    // create the builder
    IBuilder* builder = createInferBuilder(gLogger);//创建构建器(即指向Ibuilder类型对象的指针)
    IBuilderConfig *config = builder->createBuilderConfig();
    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);    //必须加不然报错
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);/*等价于*bulider.createNetwork(),通过Ibulider定义的
    名为creatNetwork()方法，创建INetworkDefinition的对象，ntework这个指针指向这个对象*/ 
 
    auto parser = nvonnxparser::createParser(*network, gLogger.getTRTLogger());//创建解析器
 
    //Optional - uncomment below lines to view network layer information
    //config->setPrintLayerInfo(true);
    //parser->reportParsingInfo();
 
    if (!parser->parseFromFile(modelFile.c_str(), verbosity)) //解析onnx文件，并填充网络
    
        string msg("failed to parse onnx file");
        gLogger.log(nvinfer1::ILogger::Severity::kERROR, msg.c_str());
        exit(EXIT_FAILURE);
    
 
    // Build the engine
    builder->setMaxBatchSize(maxBatchSize);
    config->setMaxWorkspaceSize(1 << 30);
    // builder->setMaxWorkspaceSize(1 << 30);
#ifdef USE_FP16
	config->setFlag(BuilderFlag::kFP16);
#endif
    // samplesCommon::enableDLA(builder, gUseDLACore);
    //当引擎建立起来时，TensorRT会复制
    // ICudaEngine* engine = builder->buildCudaEngine(*network);//通过Ibuilder类的buildCudaEngine()方法创建IcudaEngine对象，
    ICudaEngine *engine = builder->buildEngineWithConfig(*network,*config);
    assert(engine);
 
    // we can destroy the parser
    parser->destroy();
    
    // serialize the engine, 
    // then close everything down
    trtModelStream = engine->serialize();//将引擎序列化，保存到文件中
    engine->destroy();
    network->destroy();
    builder->destroy();
    config->destroy();
    saveToTrtModel(TrtSaveFileName,trtModelStream);



void onnxToTRTModelDynamicBatch(const std::string& modelFile, unsigned int maxBatchSize, IHostMemory*& trtModelStream,const char * TrtSaveFileName,int input_h,int input_w) // output buffer for the TensorRT model 动态batch

            int verbosity = (int) nvinfer1::ILogger::Severity::kWARNING;
        
            // create the builder
            IBuilder* builder = createInferBuilder(gLogger);//创建构建器(即指向Ibuilder类型对象的指针)
            IBuilderConfig *config = builder->createBuilderConfig();
            auto profile = builder->createOptimizationProfile();

            const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);    //必须加不然报错
            nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);/*等价于*bulider.createNetwork(),通过Ibulider定义的
            名为creatNetwork()方法，创建INetworkDefinition的对象，ntework这个指针指向这个对象*/ 

            Dims dims = Dims41, 3, input_h, input_w;
            profile->setDimensions(INPUT_BLOB_NAME,
                        OptProfileSelector::kMIN, Dims41, dims.d[1], dims.d[2], dims.d[3]);
            profile->setDimensions(INPUT_BLOB_NAME,
                        OptProfileSelector::kOPT, Dims4maxBatchSize, dims.d[1], dims.d[2], dims.d[3]);
            profile->setDimensions(INPUT_BLOB_NAME,
                        OptProfileSelector::kMAX, Dims4maxBatchSize, dims.d[1], dims.d[2], dims.d[3]);
            config->addOptimizationProfile(profile);
        
            auto parser = nvonnxparser::createParser(*network, gLogger.getTRTLogger());//创建解析器
        
            //Optional - uncomment below lines to view network layer information
            //config->setPrintLayerInfo(true);
            //parser->reportParsingInfo();
        
            if (!parser->parseFromFile(modelFile.c_str(), verbosity)) //解析onnx文件，并填充网络
            
                string msg("failed to parse onnx file");
                gLogger.log(nvinfer1::ILogger::Severity::kERROR, msg.c_str());
                exit(EXIT_FAILURE);
            
        
            // Build the engine
            // builder->setMaxBatchSize(maxBatchSize);
            config->setMaxWorkspaceSize(1 << 30);
            // builder->setMaxWorkspaceSize(1 << 30);
        #ifdef USE_FP16
            config->setFlag(BuilderFlag::kFP16);
        #endif
            // samplesCommon::enableDLA(builder, gUseDLACore);
            //当引擎建立起来时，TensorRT会复制
            // ICudaEngine* engine = builder->buildCudaEngine(*network);//通过Ibuilder类的buildCudaEngine()方法创建IcudaEngine对象，
            ICudaEngine *engine = builder->buildEngineWithConfig(*network,*config);
            assert(engine);
        
            // we can destroy the parser
            parser->destroy();
            
            // serialize the engine, 
            // then close everything down
            trtModelStream = engine->serialize();//将引擎序列化，保存到文件中
            engine->destroy();
            network->destroy();
            builder->destroy();
            config->destroy();
            saveToTrtModel(TrtSaveFileName,trtModelStream);
        


int main(int argc, char** argv)

    IHostMemory* trtModelStreamnullptr;
//    int batchSize = atoi(argv[3]);
    
    // int input_h = atoi(argv[4]);
    // int input_w=atoi(argv[5]);
//    onnxToTRTModel(argv[1],batchSize,trtModelStream,argv[2]);
//    onnxToTRTModel("../onnx_model/plate_rec.onnx",batchSize,trtModelStream,"./plate_rec.trt");
    onnxToTRTModel(onnx_path, batchSize, trtModelStream, TrtSaveFileName);
    std::cout<<"convert seccuss!"<<std::endl;

：CMakeLists.txt

cmake_minimum_required(VERSION 2.6)

project(yolo)

add_definitions(-std=c++11)

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)

include_directories($PROJECT_SOURCE_DIR/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda-11.6/include)
link_directories(/usr/local/cuda-11.6/lib64)
# tensorrt
include_directories(/home/package/TensorRT-8.2.5.1/include)
link_directories(/home/package/TensorRT-8.2.5.1/lib/)

include_directories(/home/package/TensorRT-8.2.5.1/samples/common/)
#link_directories(/home/package/TensorRT-8.2.5.1/lib/stubs/)

# opencv
find_package(OpenCV REQUIRED)
include_directories($OpenCV_INCLUDE_DIRS)


add_executable(yolo $PROJECT_SOURCE_DIR/onnx2trt.cpp)
target_link_libraries(yolo nvinfer)
target_link_libraries(yolo cudart)
target_link_libraries(yolo $OpenCV_LIBS)
target_link_libraries(yolo /home/package/TensorRT-8.2.5.1/lib/stubs/libnvonnxparser.so)

add_definitions(-O2 -pthread)