c++通过tensorRT调用模型进行推理

news/2024/10/17 16:27:58/

模型来源
算法工程师训练得到的onnx模型

c++对模型的转换
拿到onnx模型后,通过tensorRT将onnx模型转换为对应的engine模型,注意:训练用的tensorRT版本和c++调用的tensorRT版本必须一致。

如何转换:

  1. 算法工程师直接转换为.engine文件进行交付。
  2. 自己转换,进入tensorRT安装目录\bin目录下,将onnx模型拷贝到bin目录,地址栏中输入cmd回车弹出控制台窗口,然后输入转换命令,如:

trtexec --onnx=model.onnx --saveEngine=model.engine --workspace=1024 --optShapes=input:1x13x512x640 --fp16

然后回车,等待转换完成,完成后如图所示:
在这里插入图片描述
并且在bin目录下生成.engine模型文件。

c++对.engine模型文件的调用和推理
首先将tensorRT对模型的加载及推理进行封装,命名为CTensorRT.cpp,老样子贴代码:

//CTensorRT.cpp
class Logger : public nvinfer1::ILogger {void log(Severity severity, const char* msg) noexcept override {if (severity <= Severity::kWARNING)std::cout << msg << std::endl;}
};Logger logger;
class CtensorRT
{
public:CtensorRT() {}~CtensorRT() {}private:std::shared_ptr<nvinfer1::IExecutionContext> _context;std::shared_ptr<nvinfer1::ICudaEngine> _engine;nvinfer1::Dims _inputDims;nvinfer1::Dims _outputDims;
public:void cudaCheck(cudaError_t ret, std::ostream& err = std::cerr){if (ret != cudaSuccess){err << "Cuda failure: " << cudaGetErrorString(ret) << std::endl;abort();}}bool loadOnnxModel(const std::string& filepath){auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));if (!builder){return false;}const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));if (!network){return false;}auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());if (!config){return false;}auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));if (!parser){return false;}parser->parseFromFile(filepath.c_str(), static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING));std::unique_ptr<IHostMemory> plan{ builder->buildSerializedNetwork(*network, *config) };if (!plan){return false;}std::unique_ptr<IRuntime> runtime{ createInferRuntime(logger) };if (!runtime){return false;}_engine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(plan->data(), plan->size()));if (!_engine){return false;}_context = std::shared_ptr<nvinfer1::IExecutionContext>(_engine->createExecutionContext());if (!_context){return false;}int nbBindings = _engine->getNbBindings();assert(nbBindings == 2); // 输入和输出,一共是2个for (int i = 0; i < nbBindings; i++){if (_engine->bindingIsInput(i))_inputDims = _engine->getBindingDimensions(i);    // (1,3,752,752)else_outputDims = _engine->getBindingDimensions(i);}return true;}bool loadEngineModel(const std::string& filepath){std::ifstream file(filepath, std::ios::binary);if (!file.good()){return false;}std::vector<char> data;try{file.seekg(0, file.end);const auto size = file.tellg();file.seekg(0, file.beg);data.resize(size);file.read(data.data(), size);}catch (const std::exception& e){file.close();return false;}file.close();auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));_engine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(data.data(), data.size()));if (!_engine){return false;}_context = std::shared_ptr<nvinfer1::IExecutionContext>(_engine->createExecutionContext());if (!_context){return false;}int nbBindings = _engine->getNbBindings();assert(nbBindings == 2); // 输入和输出,一共是2个// 为输入和输出创建空间for (int i = 0; i < nbBindings; i++){if (_engine->bindingIsInput(i))_inputDims = _engine->getBindingDimensions(i);    //得到输入结构else_outputDims = _engine->getBindingDimensions(i);//得到输出结构}return true;}void ONNX2TensorRT(const char* ONNX_file, std::string save_ngine){// 1.创建构建器的实例nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);// 2.创建网络定义uint32_t flag = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);nvinfer1::INetworkDefinition* network = builder->createNetworkV2(flag);// 3.创建一个 ONNX 解析器来填充网络nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger);// 4.读取模型文件并处理任何错误parser->parseFromFile(ONNX_file, static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING));for (int32_t i = 0; i < parser->getNbErrors(); ++i){std::cout << parser->getError(i)->desc() << std::endl;}// 5.创建一个构建配置,指定 TensorRT 应该如何优化模型nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();// 7.指定配置后,构建引擎nvinfer1::IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);// 8.保存TensorRT模型std::ofstream p(save_ngine, std::ios::binary);p.write(reinterpret_cast<const char*>(serializedModel->data()), serializedModel->size());// 9.序列化引擎包含权重的必要副本,因此不再需要解析器、网络定义、构建器配置和构建器,可以安全地删除delete parser;delete network;delete config;delete builder;// 10.将引擎保存到磁盘,并且可以删除它被序列化到的缓冲区delete serializedModel;}uint32_t getElementSize(nvinfer1::DataType t) noexcept{switch (t){case nvinfer1::DataType::kINT32: return 4;case nvinfer1::DataType::kFLOAT: return 4;case nvinfer1::DataType::kHALF: return 2;case nvinfer1::DataType::kBOOL:case nvinfer1::DataType::kINT8: return 1;}return 0;}int64_t volume(const nvinfer1::Dims& d){return std::accumulate(d.d, d.d + d.nbDims, 1, std::multiplies<int64_t>());}bool infer(unsigned char* input, int real_input_size, cv::Mat& out_mat){tensor_custom::BufferManager buffer(_engine);cudaStream_t stream;cudaStreamCreate(&stream); // 创建异步cuda流int binds = _engine->getNbBindings();for (int i = 0; i < binds; i++){if (_engine->bindingIsInput(i)){size_t input_size;float* host_buf = static_cast<float*>(buffer.getHostBufferData(i, input_size));memcpy(host_buf, input, real_input_size);break;}}// 将输入传递到GPUbuffer.copyInputToDeviceAsync(stream);// 异步执行bool status = _context->enqueueV2(buffer.getDeviceBindngs().data(), stream, nullptr);if (!status)return false;buffer.copyOutputToHostAsync(stream);for (int i = 0; i < binds; i++){if (!_engine->bindingIsInput(i)){size_t output_size;float* tmp_out = static_cast<float*>(buffer.getHostBufferData(i, output_size));//do your something herebreak;}}cudaStreamSynchronize(stream);cudaStreamDestroy(stream);return true;}
};

调用方式

int main()
{vector<int> dims = { 1,13,512,640 };vector<float> vall;for (int i=0;i<13;i++){string file = "D:\\xxx\\" + to_string(i) + ".png";cv::Mat mt = imread(file, IMREAD_GRAYSCALE);cv::resize(mt, mt, cv::Size(640,512));mt.convertTo(mt, CV_32F, 1.0 / 255);cv::Mat shape_xr = mt.reshape(1, mt.total() * mt.channels());std::vector<float> vec_xr = mt.isContinuous() ? shape_xr : shape_xr.clone();vall.insert(vall.end(), vec_xr.begin(), vec_xr.end());}cv::Mat mt_4d(4, &dims[0], CV_32F, vall.data());string engine_model_file = "model.engine";CtensorRT cTensor;if (cTensor.loadEngineModel(engine_model_file)){cv::Mat out_mat;if (!cTensor.infer(mt_4d.data, vall.size() * 4, out_mat))std::cout << "infer error!" << endl;elsecv::imshow("out", out_mat);}elsestd::cout << "load model file failed!" << endl;cv::waitKey(0);return 0;
}

http://www.ppmy.cn/news/1095882.html

相关文章

WMS仓储管理系统的功能与WCS系统有什么区别

在物流行业的现代化管理中&#xff0c;WMS仓储管理系统和WCS仓库控制系统是两个重要的组成部分。虽然它们都是用于仓库管理的软件系统&#xff0c;但是它们的功能和应用场景有很大的区别。本文我们将详细阐述这两者的功能和区别。 一、WMS仓储管理系统 WMS是一种软件系统&…

仓储24代电子标签接口文档

电子标签标签注册 通过手动触发电子标签注册到系统&#xff0c;注册成功就可以进行功能测试。 仓储24代注册 注册成功显示rssi:-xxdB如下&#xff1a; 绝对值越小信号越差&#xff0c;一般0 ~ -75dB都可以正常通讯 查询电子标签信息接口 接口描述 查询注册到系统的电子标签的…

Spark【RDD编程(四)综合案例】

案例1-TOP N个数据的值 输入数据&#xff1a; 1,1768,50,155 2,1218,600,211 3,2239,788,242 4,3101,28,599 5,4899,290,129 6,3110,54,1201 7,4436,259,877 8,2369,7890,27 处理代码&#xff1a; def main(args: Array[String]): Unit {//创建SparkContext对象val conf…

项目01—基于nignx+keepalived双vip的负载均衡高可用Web集群

文章目录 一.项目介绍1.拓扑图2.详细介绍 二.前期准备1.项目环境2.IP划分 三. 项目步骤1.ansible部署软件环境1.1 安装ansible环境1.2 建立免密通道1.3 批量部署nginx 2.配置NFS服务器和负载均衡器搭建keepalived2.1 修改nginx的index.html界面2.2 nginx实现七层负载均衡2.4 使…

Json“牵手”当当网商品详情数据方法,当当商品详情API接口,当当API申请指南

当当网是知名的综合性网上购物商城&#xff0c;由国内著名出版机构科文公司、美国老虎基金、美国IDG集团、卢森堡剑桥集团、亚洲创业投资基金&#xff08;原名软银中国创业基金&#xff09;共同投资成立1。 当当网从1999年11月正式开通&#xff0c;已从早期的网上卖书拓展到网…

微服务05-Docker基本操作

Docker的定义 1.什么是Docker Docker是一个快速交付应用、运行应用的技术&#xff1a; 可以将程序及其依赖、运行环境一起打包为一个镜像&#xff0c;可以迁移到任意Linux操作系统运行时利用沙箱机制形成隔离容器&#xff0c;各个应用互不干扰启动、移除都可以通过一行命令完…

SpringMVC概述与简单使用

1.SpringMVC简介 SpringMVC也叫做Spring web mvc,是 Spring 框架的一部分&#xff0c;是在 Spring3.0 后发布的。 2.SpringMVC优点 1.基于 MVC 架构 基于 MVC 架构&#xff0c;功能分工明确。解耦合&#xff0c; 2.容易理解&#xff0c;上手快&#xff1b;使用简单。 就可以…

云计算时代的采集利器

大家好&#xff01;在今天的知识分享中&#xff0c;我们将探讨一个在云计算环境中的爬虫应用利器——独享IP。如果你是一名爬虫程序员&#xff0c;或者对数据采集和网络爬虫有浓厚的兴趣&#xff0c;那么这篇文章将向你展示独享IP在云计算环境下的应用价值。 1. 什么是独享IP&…