![82fd80704bf3dc821c8750a0a10ffd13.png](https://i-blog.csdnimg.cn/blog_migrate/96299e43bbeb4b56fb50608ec8baca65.png)
简介
- tensorflow、pytorch、caffe等训练框架训练好的模型不能直接用于tensorrt推理,需要转为tensorrt需要模型格式engine,主流有caffe和onnx转engine, 这里介绍onnx转engine, 因为训练模型转为onnx比较容易。
具体步骤
定义网络->加载onnx模型和配置信息-->转为engine文件。
- 0)init logger
#include <iostream>
#include "NvInferRuntimeCommon.h"
class Logger: public nvinfer1::ILogger
{
void log(Severity severity, const char* msg) override
{
// suppress info-level messages
if (severity != Severity::kINFO)
std::cout << msg << std::endl;
}
}gLogger;
- 1) create builder
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
- 2) create network
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U
<<static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
- 3) create parser
nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
// onnx 模型
parser->parseFromFile("model.onnx", int(nvinfer1::ILogger::Severity::kWARNING));
// 打印相关信息
for (int i = 0; i < parser->getNbErrors(); ++i)
{
std::cout << parser->getError(i)->desc() << std::endl;
}
- 4) create config
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(1<<20);
- 5) create engine and save engine model,输出model.engine文件,用于tensorRT推理。
nvinfer1::ICudaEngine* engine = build->buildEngineWithConfig(*network, *config);
// serialize
nvinfer1::IHostMemory *seriallizedModel = engine->serialize();
// save engine
std::ofstream p("model.engine", std::ios::binary);
if (!p)
{
std::cerr << "could not open plan output file" << std::endl;
return false;
}
p.write(reinterpret_cast<const char*>(seriallizedModel->data()), seriallizedModel->size());