c++通过onnxruntime调用sklearn

最新推荐文章于 2024-09-10 08:16:42 发布

欧拉欧拉木大

最新推荐文章于 2024-09-10 08:16:42 发布

阅读量1.6k

点赞数 4

分类专栏： c++ 文章标签： sklearn c++ python

本文链接：https://blog.csdn.net/weixin_42256401/article/details/130319752

版权

c++ 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

文章介绍了如何将Python的sklearn训练的随机森林模型转换为ONNX格式，并在C++中使用ONNXRuntime进行推理。首先，通过skl2onnx库将模型转换为ONNX模型并保存，然后在C++中加载模型，获取输入输出节点信息，并创建输入张量进行预测。

摘要由CSDN通过智能技术生成

概述

python的sklearn很方便训练一些机器学习模型，但是c++调用推理时会出现各种意料之外的问题。可以实现c++用onnxruntime调用sklearn。

训练sklearn并将模型保存为onnx

以一个训练好的随机森林模型为例

import onnxmltools
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import *
import pickle
with open("rfc.pkl", "rb") as f:
    model = pickle.load(f)
#定义输入数据的type和shape
initial_type = [('float_input',DoubleTensorType([None, 5]))]
onnx_model = convert_sklearn(model,initial_types=initial_type)
onnxmltools.utils.save_model(onnx_model, "model.onnx")

import onnxruntime as rt
import numpy as np
sess = rt.InferenceSession("model.onnx")
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
probability_name = sess.get_outputs()[1].name

feature = data # load data
pred_onx = sess.run([label_name, probability_name], {input_name: feature})

c++调用

#include <iostream>
#include <vector>
#include <assert.h>
#include <core/session/onnxruntime_cxx_api.h>
#include <core/providers/cuda/cuda_provider_factory.h>
#include <core/session/onnxruntime_c_api.h>
int main(int argc,char** argv)
{
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "Default");
    #ifdef _WIN32
        const wchar_t* model_path = L"model.onnx";
    #else
        const char* model_path = "model.onnx";
    #endif
    Ort::Session session{env, model_path, Ort::SessionOptions{nullptr}}; //CPU
    Ort::AllocatorWithDefaultOptions allocator;
    
	// 获取模型输入输出的节点个数
    size_t num_input_nodes = session.GetInputCount();
    size_t num_output_nodes = session.GetOutputCount();
    
    // 打印模型中输入和输出的名称
    for(int i = 0; i < num_input_nodes; i++)
    {
        char* in_name = session.GetInputName(i,allocator);
        std::cout<<in_name<<std::endl;
    }
    for(int i = 0; i < num_output_nodes; i++)
    {
        char* out_name = session.GetOutputName(i,allocator);
        std::cout<<out_name<<std::endl;
    }
    // 输出输出节点必须与上面模型的输入输出节点名一致
    std::vector<const char*> input_node_names = {"float_input"};
    std::vector<const char*> output_node_names = {"output_label","output_probability"};
    // 设置输出数据的维度，这里以单条数据为例
    std::vector<int64_t> input_node_dims = {1, 5};
    size_t input_tensor_size = 1 * 5; 
    // 构造输入数据
    std::vector<double> input_tensor_values(input_tensor_size);
    for (unsigned int i = 0; i < input_tensor_size; i++)
    {
        input_tensor_values[i] = (double)i / (input_tensor_size + 1);
        std::cout<<input_tensor_values[i]<<std::endl;
    }
    
    // create input tensor object from data values
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<double>(memory_info, input_tensor_values.data(), input_tensor_size, input_node_dims.data(), 2);
    assert(input_tensor.IsTensor());

    std::vector<Ort::Value> ort_inputs;
    ort_inputs.push_back(std::move(input_tensor));
    // score model & input tensor, get back output tensor
    auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_node_names.data(), ort_inputs.data(), ort_inputs.size(), output_node_names.data(), 2);

    // Get pointer to output tensor float values
    int* floatarr = output_tensors[0].GetTensorMutableData<int>();
    std::cout<<*floatarr<<std::endl;
}