训练好的模型（如.pt）转成onnx形式，ONNX定义了一组与环境和平台无关的标准格式。ONNX文件不仅存储了神经网络模型的权重，还存储了模型的结构信息、网络中各层的输入输出等一些信息。
ONNX的推理可以用ONNX Runtime官方库，如果在英伟达平台上，可以转TensorRT后运行。本文主要介绍转TRT格式后如何C++部署运行。
1、ONNX 转 RTR
基本流程就是：
1、创建构建器，由构建器创建网络，然后解析器解析ONNX文件。
2、设置一些必要参数
3、构建其构建网络然后保存成TRT模型
用到的logging.h 文件直接用NVIDIA自带的。
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "logging.h"
#include<fstream>
#include<iostream>
#include<string>
using namespace std;
using namespace nvonnxparser;
using namespace nvinfer1;
#define USE_FP16
static Logger gLogger;
void saveToTrtModel(std::string trt_save_path,IHostMemory*trtModelStream){
std::ofstream out(trt_save_path, std::ios::binary);
if (!out.is_open()){
std::cout << "打开文件失败!" <<std:: endl;
}
out.write(reinterpret_cast<const char*>(trtModelStream->data()), trtModelStream->size());
out.close();
}
int onnx2trt(){
std::string onnx_path = "../onnx_model/plate_detect.onnx";
std::string trt_save_path = "../onnx_model/plate_detect.trt";
int batch_size = 1;
IBuilder * builder = createInferBuilder(gLogger);
INetworkDefinition *network = builder->createNetworkV2(1U);
// 解析模型
IParser *parser = nvonnxparser::createParser(*network, gLogger);
if(!parser->parseFromFile(onnx_path.c_str(), (int)nvinfer1::ILogger::Severity::kWARNING)){
std::cout << " parse onnx file fail ..." << std::endl;
return -1;
}
IBuilderConfig *config = builder->createBuilderConfig();
builder->setMaxBatchSize(batch_size);
config->setMaxWorkspaceSize(1<<30);
auto profile = builder->createOptimizationProfile();
auto input_tensor=network->getInput(0);
auto input_dims = input_tensor->getDimensions();
input_dims.d[0] = 1;
profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, input_dims);
profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, input_dims);
input_dims.d[0] = batch_size;
profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);
config->addOptimizationProfile(profile);
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
#ifdef USE_INT8
config->setFlag(BuilderFlag::kINT8);
#endif
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
assert(engine);
IHostMemory* trtModelStream = engine->serialize(); //序列化 保存trt
saveToTrtModel(trt_save_path.c_str(), trtModelStream);
parser->destroy();
engine->destroy();
network->destroy();
config->destroy();
builder->destroy();
return 0;
}
cmake_minimum_required(VERSION 3.10)
project(onnx2trt)
include_directories(/home/max/TensorRT-7.2.1.6/include)
link_directories(/home/max/TensorRT-7.2.1.6/lib)
add_executable(onnx2trt onnx2trt.cpp)
target_link_libraries(onnx2trt nvinfer)
target_link_libraries(onnx2trt cudart)
target_link_libraries(onnx2trt nvonnxparser)
2、 RTR 模型推理
转好的TRT模型在部署工程上推理运行。
基本流程：
1、trt从文件中解析出模型，并反序列化到推理CUDA推理引擎。
2、分配推理所需要的CPU、GPU内存空间
3、引擎推理获取结果
调用以下类接口即可推理。
#include "logging.h"
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include <NvInferRuntime.h>
#include <cuda_runtime.h> // cuda include
#include<fstream>
#include<iostream>
#include<string>
using namespace nvinfer1;
static Logger gLogger;
class TrtDetct
{
private:
char *_trtModelStream{
nullptr};
IRuntime* _runtime = nullptr;
ICudaEngine* _engine=nullptr;
IExecutionContext* _context=nullptr;
void *_inferbuffers[2];
int _outputSize = 0;
int _input_h = 640;
int _input_w = 640;
cudaStream_t _stream;
private:
int getoutputSize(){
auto out_dims = _engine->getBindingDimensions(1);
int outputSize = 1;
for(int j = 0; j < out_dims.nbDims; j++) {
std::cout << "j = " << j << " size = " << out_dims.d[j] << std::endl;
outputSize *= out_dims.d[j];
}
return outputSize;
}
public:
TrtDetct(/* args */){
};
~TrtDetct(){
if (nullptr != _trtModelStream){
delete [] _trtModelStream;
}
};
// 文件读取模型，并反序列化成engine
void load_trtmodel(std::string trt_model_path){
std::ifstream file(trt_model_path, std::ios::binary);
size_t size{
0};
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
_trtModelStream = new char[size];
assert(_trtModelStream);
file.read(_trtModelStream, size);
file.close();
}
_runtime = createInferRuntime(gLogger);
assert(_runtime != nullptr);
_engine = _runtime->deserializeCudaEngine(_trtModelStream, size);
assert(_engine != nullptr);
_context = _engine->createExecutionContext();
assert(_context != nullptr);
initbuff();
}
//分配处理相关内存
void initbuff(){
_outputSize = getoutputSize();
// 这两个值在生成onnx时刻已经固定
const int inputIndex = _engine->getBindingIndex("input");
const int outputIndex = _engine->getBindingIndex("output");
assert(inputIndex == 0);
assert(outputIndex == 1);
CHECK(cudaMalloc((void**)&_inferbuffers[inputIndex], 3 * _input_h * _input_w * sizeof(float))); //trt输入内存申请
CHECK(cudaMalloc((void**)&_inferbuffers[outputIndex], _outputSize * sizeof(float))); //trt输出内存申请
CHECK(cudaStreamCreate(&_stream));
}
// 推理
void infer_trtmodel(){
//图像数据填充_inferbuffers[0],GPU CUDA处理
_context->enqueueV2((void **)_inferbuffers, _stream, nullptr);
//_inferbuffers[1]模型输出后处理,可以GPU处理，否则拷贝到cpu处理
}
};
CMakeLists.txt 文件如下：
project(trt_detect)
add_definitions(-std=c++11)
add_definitions(-w)
find_package(CUDA REQUIRED)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Release)
#cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
include_directories(/home/max/TensorRT-7.2.1.6/include)
link_directories(/home/max/TensorRT-7.2.1.6/lib)
cuda_add_executable(trt_detect trt_detect.cpp)
target_link_libraries(trt_detect nvinfer)
target_link_libraries(trt_detect cudart)
target_link_libraries(trt_detect nvonnxparser)
add_definitions(-O2)