ONNX Runtime C++端侧模型部署YOLOv5

加载准备

初始化ONNXRuntime环境

Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "YOLOv5Inference");

Ort::Env 是 ONNX Runtime C++ API 中用于初始化运行环境的类,有多个重载的构造函数,下面是一个构造函数原型及参数作用如下。

Ort::Env(
    OrtLoggingLevel logging_level,
    const char* logid,
    OrtLoggingFunction logging_fn = nullptr,
    void* logger_param = nullptr
);

  • logging_level:控制日志输出级别
  • logid: 自定义日志标签,用于区分不同模块的日志来源
  • logging_fn:自定义日志回调函数,若为 nullptr 则使用默认日志输出到控制台。
  • logger_param:传递给自定义日志函数的用户参数(如上下文对象)

设置会话参数

Ort::SessionOptions session_options;

session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

初始化一个空的会话配置对象session_options,SetIntraOpNumThreads限制单个算子(Intra-op)内部使用的线程数为 1,适用于轻量级任务或避免多线程竞争,SetGraphOptimizationLevel启用所有图优化策略(如算子融合、常量折叠),提升推理性能。

模型加载

Ort::Session session_(env, modelPath.c_str(), session_options);

Ort::Session 是 ONNX Runtime C++ API 中用于加载 ONNX 模型并创建推理会话的核心类,其功能分解如下。

Ort::Session(
    const Ort::Env& env,
    const char* model_path,
    const Ort::SessionOptions& options
);

  • env:全局运行环境对象,管理线程池和内存分配等资源,需优先初始化。
  • model_path:ONNX 模型文件的路径,c语言的字符串类型。
  • options:会话参数,配置会话行为,如线程数、优化级别、硬件后端等。

获取输入和输出信息

输入名称

    Ort::AllocatorWithDefaultOptions allocator;
    //创建默认内存分配器对象,用于管理 ONNX Runtime 中的内存分配(如节点名称字符串的内存
    std::vector<const char*> input_node_names_;
    //存储 C 风格字符串指针,用于直接传递给 ONNX Runtime 的推理接口
    std::vector<std::string> input_names_;
    //存储标准字符串对象的vector,用于长期维护字符串内存

    size_t num_inputs_;
    num_inputs_ = session_.GetInputCount();
    //获取输入节点的个数,有多少个节点就决定了多个个name,一般都是1个。
    input_node_names_.resize(num_inputs_);
    input_names_.resize(num_inputs_, "");
    //预分配容器空间,避免动态扩容的开销
    std::cout << "num_inputs = "<< num_inputs_<<std::endl;
    for (size_t i = 0; i < num_inputs_; ++i) {
        auto input_name = session_.GetInputNameAllocated(i, allocator);
        //通过分配器安全获取第 i 个输入节点的名称(返回 Ort::AllocatedStringPtr 对象)
        input_names_[i].append(input_name.get());
        //获取名称的原始指针,存入 input_names_ 的字符串中
        input_node_names_[i] = input_names_[i].c_str();
        //将 std::string 转换为 C 风格指针,供 input_node_names_ 使用
    }

上面函数示例了如何获取输入节点name,首先通过session_.GetInputCount()获取到输入的节点,然后使用for循环进行遍历每个节点,通过session_.GetInputNameAllocated(i, allocator)获取每个节点的名称,返回一个Ort::AllocatedStringPtr智能指针,需要通过.get方法返回c字符串,由于智能指针指向的存储空间退出for后会销毁,所以上述代码将其复制到input_names_中。

输入张量维度

    Ort::TypeInfo input_type_info = session_.GetInputTypeInfo(0);
    //获取模型第0个输入节点的类型信息对象,返回Ort::TypeInfo类型
    auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
    //从类型信息中提取张量相关的形状和数据类型信息,返回Ort::TensorTypeAndShapeInfo对象
    std::vector<int64_t> input_dims = input_tensor_info.GetShape();
    //获取输入张量的维度信息,返回std::vector<int64_t>容器,存储各维度大小
    //典型YOLO模型的输入维度为[batch, channel, height, width]
    int inputWidth = input_dims[3];
    int inputHeight = input_dims[2];

上面函数示例获取输入张量形状,最终通过张量的形状获取到了输入图像的宽和高。实际上可以简化一下,按照下面的方式。

auto inputShapeInfo = session_.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
int ch = inputShapeInfo[1];
inputWidth = inputShapeInfo[2];
inputHeight = inputShapeInfo[3];

输出名称

    std::vector<const char*> output_node_names_;
    //存储C风格字符串指针的vector,用于兼容需要const char*的ONNX Runtime API调用
    std::vector<std::string> output_names_;
    //存储标准字符串对象的vector,用于长期维护字符串内存
    size_t num_outputs_;
    num_outputs_ = session_.GetOutputCount();
    //获取模型输出节点数量
    output_node_names_.resize(num_outputs_);
    output_names_.resize(num_outputs_, "");
    //预分配两个vector的空间
    for (size_t i = 0; i < num_outputs_; ++i) {
        auto output_name = session_.GetOutputNameAllocated(i, allocator);
        output_names_[i].append(output_name.get());
        //将名称存入std::string保证生命周期
        output_node_names_[i] = output_names_[i].c_str();
    }
    //循环获取每个输出节点名称

上面示例了获取输出名称,与输入方法类似。

输入预处理

    cv::Mat image = cv::imread(imagePath);
    if (image.empty()) {
        std::cerr << "Error: Could not read image." << std::endl;
        return -1;
    }
    cv::Mat originalImage = image.clone();
    cv::Size image_shape = originalImage.size();

    // 图像预处理
    std::vector<float> inputTensor = preprocess(image, inputWidth, inputHeight);

使用opencv读取图像,调用preprocess进行预处理。

std::vector<float> preprocess(const cv::Mat& image, int inputWidth = 320, int inputHeight = 320) {
    cv::Mat resizedImage;
    cv::resize(image, resizedImage, cv::Size(inputWidth, inputHeight));
    //图像缩放:使用OpenCV的resize函数将图像调整为指定尺寸(默认320x320)
    cv::cvtColor(resizedImage, resizedImage, cv::COLOR_BGR2RGB);
    //颜色空间转换:从BGR转换为RGB格式(多数深度学习模型使用RGB输入)
    resizedImage.convertTo(resizedImage, CV_32F, 1.0 / 255.0);
    //数值归一化:通过convertTo将像素值从[0,255]归一化到[0,1]范围

    std::vector<float> inputTensor;
    for (int c = 0; c < 3; ++c) {
        for (int h = 0; h < inputHeight; ++h) {
            for (int w = 0; w < inputWidth; ++w) {
                inputTensor.push_back(resizedImage.at<cv::Vec3f>(h, w)[c]);
            }
        }
    }
    //通过三重循环将OpenCV的HWC格式(Height-Width-Channel)转换为CHW格式
    //内存布局变为连续通道数据:RRR...GGG...BBB,最终输出std::vector<float>
    return inputTensor;
}

上面的代码实现了模型对输入数据的部分预处理,包括输入图片缩放固定尺寸,数值归一化,以及将格式转换为CHW张量格式,但是对于输入模型,需要的数据格式为Ort::Value类型。

std::vector<int64_t> input_shape = {1, 3, inputHeight, inputWidth};
//input_shape采用NCHW格式(批次数-通道-高度-宽度),这是深度学习模型的通用输入布局
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
//描述用于描述内存分配的信息,包括内存的位置(CPU 或 GPU)以及内存的具体类型(固定内存或常规内存)

Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info,
        inputTensor.data(), inputTensor.size(),
        input_shape.data(), input_shape.size());

关于CreateTensor对象解析如下。

static Ort::Value CreateTensor<float>(
    const MemoryInfo& memory_info,  // 内存管理策略
    float* p_data,                  // 输入数据指针
    size_t p_data_length,           // 输入的大小
    const int64_t* shape,           // 维度数组指针
    size_t shape_length             // 维度数量
);

  • memory_info:指定张量内存分配策略,通常由Ort::MemoryInfo::CreateCpu创建。
  • inputTensor.data():输入数据的地址(需确保内存连续)。
  • inputTensor.size():输入数据的大小。
  • input_shape.data():输入张量的形状数组信息(如NCHW格式的{1,3,640,640})。
  • shape_length: 输入张量的形状信息维度数

模型推理

std::vector<Ort::Value> outputs = session_.Run(
    Ort::RunOptions{nullptr},
    input_node_names_.data(),
    &input_tensor,
    1,
    output_node_names_.data(),
    output_node_names_.size());

下面是函数的参数

OrtStatus * OrtApi::Run(
    const OrtRunOptions * run_options,
    const char *const *input_names,  //输入节点名称的数组
    const OrtValue *const *inputs,   //模型输入的数据Ort::Value类型
    size_t  input_len,  //输入张量数量,需与input_names数组长度一致
    const char *const *output_names,//输出节点名称数组
    size_t  output_names_len,//输出节点名称的数量,与output_names数组数量保持一致。
)

输出后处理

    //张量信息提取,outputs[0]指向坐标、分数张量指针,outputs[1]指向类别张量的指针
    float* dets_data = outputs[0].GetTensorMutableData<float>();
    //坐标(格式为[x1,y1,x2,y2,score])
    float* labels_pred_data = outputs[1].GetTensorMutableData<float>();
    //类别

    //张量维度的解析,用于获取检测框的数量
    auto dets_tensor_info = outputs[0].GetTensorTypeAndShapeInfo();
    std::vector<int64_t> dets_dims = dets_tensor_info.GetShape();
    size_t num_detections = dets_dims[1];

    //结构化重组,解析输出的张量将其存储dets、scores、lables_pred
    std::vector<std::vector<float>> dets(num_detections, std::vector<float>(4));
    std::vector<float> scores(num_detections);
    std::vector<int> labels_pred(num_detections);

    //遍历解析存储坐标dets、分数scores、标签类别lables_pred
    for (size_t i = 0; i < num_detections; ++i) {
        for (int j = 0; j < 4; ++j) {
            dets[i][j] = dets_data[i * 5 + j];
        }
        scores[i] = dets_data[i * 5 + 4];
        labels_pred[i] = static_cast<int>(labels_pred_data[i]);
    }

    //将坐标信息进行缩放以适应正常的图片大小。
    float scale_x = static_cast<float>(image_shape.width) / inputWidth;
    float scale_y = static_cast<float>(image_shape.height) / inputHeight;
    for (auto& det : dets) {
        det[0] *= scale_x;
        det[1] *= scale_y;
        det[2] *= scale_x;
        det[3] *= scale_y;
    }

上面的代码从输出张量信息中进行解析,将坐标、分数、标签类别依次存储到dets、scores、lables_pred中。

void visualizeResults(cv::Mat& image, const std::vector<std::vector<float>>& dets,
    const std::vector<float>& scores,
    const std::vector<int>& labels_pred,
    const std::vector<std::string>& labels,
    float conf_threshold = 0.4) {

    for (size_t i = 0; i < dets.size(); ++i) {
        const auto& det = dets[i];
        float score = scores[i];
        if (score > conf_threshold) {
            int class_id = labels_pred[i];
            int x1 = static_cast<int>(det[0]);
            int y1 = static_cast<int>(det[1]);
            int x2 = static_cast<int>(det[2]);
            int y2 = static_cast<int>(det[3]);
            std::string label = labels[class_id];
            cv::rectangle(image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(0, 255, 0), 2);
            cv::putText(image, label + ": " + std::to_string(score), cv::Point(x1, y1 - 10), cv::FONT_HERSHEY_SIMPLEX, 0.9, cv::Scalar(0, 255, 0), 2);
        }
    }
}

最终将获取到的将坐标、分数、标签类别传入到visualizeResults进行绘制。

发现一个开源的ai toolkit,相对比较全。https://github.com/xlite-dev/lite.ai.toolkit/tree/main