ONNX - YOLO 神经网络模型的部署攻略

2023-09-09 01:52:51

在 C++ 中部署 YOLO 模型到 ONNX 格式：实现目标检测

跨平台推理：为什么选择 ONNX？

ONNX（开放神经网络交换） 是一种模型格式，可实现神经网络模型在不同框架和平台之间的轻松交换和部署。其强大的跨平台兼容性使其成为将模型部署到各种系统中的理想选择。ONNX 还以其灵活性、可扩展性和经过优化的性能而闻名，确保高性能推理速度和准确性。

步骤指南：部署 YOLO 模型到 ONNX

1. 下载项目

从 bubbliiiing 的 GitHub 存储库下载示例项目，其中包含将 YOLO 模型导出到 ONNX 格式的代码。

2. 准备导出函数

在下载的项目中，将 predict 函数重命名为 export_onnx 函数，使其负责导出 YOLO 模型到 ONNX 格式。

3. 指定 YOLO 模型版本

编辑文件夹 nets 中的 yolo.py 文件，指定要导出的 YOLO 模型版本。

4. 生成 ONNX 文件

运行项目以生成 ONNX 文件。

5. 编译 ONNX 文件

下载另一个项目以对 ONNX 文件执行 cmake 和 make 操作。

6. 加载和执行导出的 ONNX 模型

使用 C++ 代码加载和执行导出的 ONNX 模型，实现目标检测。

代码示例：加载 ONNX 模型

#include <iostream>
#include <vector>

#include "onnxruntime.h"

int main() {
  // 1. 加载 ONNX 模型
  Ort::SessionOptions session_options;
  Ort::Session session(session_options, "yolov3.onnx");

  // 2. 设置输入张量
  std::vector<float> input_data = { ... }; // 准备图像数据
  Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
      session, input_data.data(), input_data.size(), input_shape.data(), input_shape.size());

  // 3. 运行推理
  std::vector<Ort::Value> output_tensors = session.Run(Ort::RunOptions(), {input_tensor});

  // 4. 处理推理结果
  // ...

  return 0;
}