swift int 转 data_TensorRT | onnx模型转engine模型

简介

• 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());