#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 enginefilePath, int type) {
// 创建builder
IBuilder* builder = createInferBuilder(gLogger);
// 创建network
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"); // 解析模型成功,第一个断点测试
// 使用builder对象构建engine
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(16 * (1 << 20)); // 设置最大工作空间
config->setFlag(BuilderFlag::kGPU_FALLBACK); // 启用GPU回退模式,作用?
// config->setFlag(BuilderFlag::kSTRICT_TYPES); //强制执行xx位的精度计算
if (type == 1) {
config->setFlag(nvinfer1::BuilderFlag::kFP16); // 设置精度计算
}
if (type == 2) {
config->setFlag(nvinfer1::BuilderFlag::kINT8);
}
IOptimizationProfile* profile = builder->createOptimizationProfile(); //创建优化配置文件
profile->setDimensions("x", OptProfileSelector::kMIN, Dims4(1, 3, 32, 300)); // 设置输入x的动态维度,最小值
profile->setDimensions("x", OptProfileSelector::kOPT, Dims4(1, 3, 32, 320)); // 期望输入的最优值
profile->setDimensions("x", OptProfileSelector::kMAX, Dims4(1, 3, 32, 340)); // 最大值
config->addOptimizationProfile(profile);
ICudaEngine* myengine = builder->buildEngineWithConfig(*network, *config); //创建engine 第二个断点测试
std::cout << "try to save engine file now" << std::endl;
std::ofstream p(enginefilePath, std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return 0;
}
// 序列化
IHostMemory* modelStream = myengine->serialize(); // 第三个断点测试
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size()); // 写入
modelStream->destroy(); // 销毁
myengine->destroy();
network->destroy();
parser->destroy();
std::cout << "convert onnx model to TensorRT engine model successfully!" << std::endl; // 转换成功,第四个断点测试
return 0;
}
int main(int argc, char** argv) {
onnx2engine("yolov5s.onnx", "yolov5s.engine", 1);
return 0;
}
onnx转成engine C++实现
最新推荐文章于 2024-05-28 14:47:52 发布