C++ 上用 ONNXruntime 部署自己的模型

原创

已于 2022-02-21 14:35:10 修改 · 2.5w 阅读

183 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #tensorflow #onnxruntime #模型部署

于 2022-01-21 17:56:53 首次发布

本文介绍如何使用C++结合ONNXruntime部署一个基于Keras搭建的相机位姿估计网络模型，并利用TensorRT加速推理过程。文章详细展示了模型准备、ONNXruntime配置及模型部署的具体步骤。

利用C++ ONNXruntime部署自己的模型，这里用Keras搭建好的一个网络模型来举例，转换为onnx的文件，在C++上进行部署，另外可以利用tensorRT加速。

Github地址：https://github.com/zouyuelin/SLAM_Learning_notes/tree/main/PoseEstimation

网盘地址
链接：https://pan.baidu.com/s/19ncKS8HhwDaYxe2WGlUmNQ?pwd=car0
提取码：car0

一、模型的准备

搭建网络模型训练：
tensorflow keras 搭建相机位姿估计网络–例
网络的输入输出为：
网络的输入： [image_ref , image_cur]
网络的输出: [tx , ty , tz , roll , pitch , yaw]

训练的模型位置：kerasTempModel\，一定要用model.save()的方式，不能用model.save_model()
在onnxruntime调用需要onnx模型，这里需要将keras的模型转换为onnx模型；

安装转换的工具：

pip install tf2onnx

安装完后运行：

python -m tf2onnx.convert --saved-model kerasTempModel --output "model.onnx" --opset 14

tip:这里设置 opset 版本为14 的优化效率目前亲测是最好的，推理速度比版本 11 、12更快。

运行完以后在终端最后会告诉你网络模型的输入和输出：

2022-01-21 15:48:00,766 - INFO - 
2022-01-21 15:48:00,766 - INFO - Successfully converted TensorFlow model kerasTempModel to ONNX
2022-01-21 15:48:00,766 - INFO - Model inputs: ['input1', 'input2']
2022-01-21 15:48:00,766 - INFO - Model outputs: ['Output']
2022-01-21 15:48:00,766 - INFO - ONNX model is saved at model.onnx

模型有两个输入，输入节点名分别为['input1', 'input2']，输出节点名为['Output']。
当然也可以不用具体知道节点名，在onnxruntime中可以通过打印来查看模型的输入输出节点名。

二、配置ONNXruntime

CMakeLists.txt:
首先需要设置你的ONNXRUNTIME 的安装位置：

#******onnxruntime*****
set(ONNXRUNTIME_ROOT_PATH /path to your onnxruntime-master)
set(ONNXRUNTIME_INCLUDE_DIRS ${
   
   ONNXRUNTIME_ROOT_PATH}/include/onnxruntime
                             ${
   
   ONNXRUNTIME_ROOT_PATH}/onnxruntime
                             ${
   
   ONNXRUNTIME_ROOT_PATH}/include/onnxruntime/core/session/)
set(ONNXRUNTIME_LIB ${
   
   ONNXRUNTIME_ROOT_PATH}/build/Linux/Release/libonnxruntime.so)

C++ main.cpp中
头文件：

#include <core/session/onnxruntime_cxx_api.h>
#include <core/providers/cuda/cuda_provider_factory.h>
#include <core/session/onnxruntime_c_api.h>
#include <core/providers/tensorrt/tensorrt_provider_factory.h>

三、模型的部署

1. 模型的初始化设置

	//模型位置
    string model_path = "../model.onnx";
	//初始化设置ONNXRUNTIME 的环境
    Ort::Env env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "PoseEstimate");
    Ort::SessionOptions session_options;
    //TensorRT加速开启，CUDA加速开启
    OrtSessionOptionsAppendExecutionProvider_Tensorrt(session_options, 0); //tensorRT
    OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    Ort::AllocatorWithDefaultOptions allocator;
    //加载ONNX模型
    Ort::Session session(env, model_path.c_str(), session_options);
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

打印模型信息：printModelInfo函数

void printModelInfo(Ort::Session &session, Ort::AllocatorWithDefaultOptions &allocator)
{
   
   
    //输出模型输入节点的数量
    size_t num_input_nodes = session.GetInputCount();
    size_t num_output_nodes = session.GetOutputCount();
    cout<<"Number of input node is:"<<num_input_nodes<<endl;
    cout<<"Number of output node is:"<<num_output_nodes<<endl;

    //获取输入输出维度
    for(auto i = 0; i<num_input_nodes;i++)
    {
   
   
        std::vector<int64_t> input_dims = session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        cout<<endl<<"input "<<i<<" dim is: ";
        for(auto j=0; j<input_dims.size();j++)
            cout<<input_dims[j]<<" ";
    }
    for(auto i = 0; i<num_output_nodes;i++)
    {
   
   
        std::vector<int64_t> output_dims = session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        cout<<endl<<"output "<<i<<" dim is: ";
        for(auto j=0; j<output_dims.size();j++)
            cout<<output_dims[j]<<" ";
    }
    //输入输出的节点名
    cout<<endl;//换行输出
    for(auto i = 0; i<num_input_nodes;i++)
        cout<<"The input op-name "<<i<<" is:"<<session.GetInputName(i, allocator)<<endl;
    for(auto i = 0; i<num_output_nodes;i++)
        cout<<"The output op-name "<<i<<" is:"<<session.GetOutputName(i, allocator)<<endl;

    //input_dims_2[0] = input_dims_1[0] = output_dims[0] = 1;//batch size = 1
}

函数应用:

//打印模型的信息
printModelInfo(session,allocator);

输出结果：

Number of input node is:2
Number of output node is:1

input 0 dim is: -1 512 512 3 
input 1 dim is: -1 512 512 3 
output 0 dim is: -1 6 
The input op-name 0 is:input1
The input op-name 1 is:input2
The output op-name 0 is:Output

如果事先不知道网络，通过打印信息这时候就可以定义全局变量：

//输入网络的维度
static constexpr const int width = 512;
static constexpr const int height = 512;
static constexpr const int channel = 3;
std::array<int64_t, 4> input_shape_{
   
    1,height, width,channel};

2. 构建推理

构建推理函数computPoseDNN()步骤：

对应用Opencv输入的Mat图像进行resize：

    Mat Input_1,Input_2;
    resize(img_1,Input_1,Size(512,512));
    resize(img_2,Input_2,Size(512,512));

指定input和output的节点名，当然也可以定义在全局变量中，这里为了方便置入函数中

    std::vector<const char*> input_node_names = {
   
   "input1","input2"};
    std::vector<const char*> output_node_names = {
   
   "Output"};

分配image_ref和image_cur的内存，用指针数组存储，这里长度为 512 * 512 * 3，因为不能直接把Mat矩阵输入，所以需要数组来存储图像数据，然后再转ONNXRUNTIME专有的tensor类型即可：

    std::array<float, width * height *channel> input_image_1{
   
   };
    std::array<float, width * height *channel> input_image_2{
   
   };
    float* input_1 =  input_image_1.data();
    float* input_2 =  input_image_2