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