21xrx.com
2025-07-15 21:31:01 Tuesday
文章检索 我的文章 写文章
如何使用C++调用ONNX模型?
2023-06-28 17:36:54 深夜i     129     0
C++ ONNX模型 调用

ONNX(Open Neural Network Exchange)是一种用于表示机器学习模型的开放格式。与其他框架相比,ONNX具有跨平台和跨框架的能力,可在不同的深度学习框架中使用。

当我们有一个训练好的ONNX模型时,我们可以使用C++调用它。在下面的内容中,我们将介绍如何使用C++来加载和运行ONNX模型。

首先,我们需要安装ONNX Runtime C++库。这可以通过在终端中运行以下命令来完成:

sudo apt-get update
sudo apt-get install libonnxruntime-dev

接下来,我们需要使用ONNX Runtime C++ API来加载和评估我们的模型。我们将按照以下步骤进行:

1. 创建ONNX Runtime的推理会话(inference session)

2. 加载ONNX模型

3. 创建输入和输出向量

4. 评估模型

5. 获取输出

以下是一个通过ONNX模型实现图像分类的示例代码:

#include <onnxruntime_cxx_api.h>
#include <vector>
using namespace std;
using namespace onnxruntime;
int main() {
  // 创建inference session
  Ort::SessionOptions session_options;
  Ort::Env env;
  Ort::Session session(env, "model.onnx", session_options);
  // 加载模型
  auto input_count = session.GetInputCount();
  // 创建输入列表和输出列表
  vector<const char*> input_names(input_count);
  vector<const char*> output_names(1);
  vector<Ort::Value> input_values(input_count);
  vector<float> output_values;
  for (size_t i = 0; i < input_count; ++i) {
    // 获取输入名字和形状
    char* input_name = session.GetInputName(i, allocator);
    input_names[i] = input_name;
    Ort::TypeInfo type_info = session.GetInputTypeInfo(i);
    auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
    vector<int64_t> input_shape = tensor_info.GetShape();
    // 创建输入张量
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    auto input_array = Ort::Value::CreateTensor<float>(memory_info, input_data.data(), input_data.size(), input_shape.data(), input_shape.size());
    input_values[i] = input_array;
  }
  // 获取输出名字
  output_names[0] = session.GetOutputName(0, allocator);
  // 评估模型
  auto output_tensors = session.Run(Ort::RunOptions{}, input_names.data(), input_values.data(), input_count, output_names.data(), output_names.size());
  // 获取输出向量
  auto output_tensor = output_tensors.front().GetTensor<float>();
  auto* output_data = output_tensor.Data();
  output_values.assign(output_data, output_data + output_tensor.Shape().Size());
}

在代码中,我们首先创建了一个推理会话对象。然后,我们使用该会话对象加载ONNX模型,并创建输入和输出向量。接下来,我们将输入张量传递给模型,并获得模型的输出。最后,我们将输出向量作为结果传递给输出。

这是一个简单的实例,它展示了如何使用C++加载和评估ONNX模型。注意,该代码仅供参考,您可能需要根据您的具体情况对其进行修改。当然,您也可以使用其他方法来调用ONNX模型。总而言之,ONNX是值得使用的模型格式,它提供了开放和跨平台的方法来转移深度学习模型。

  
  

评论区