ONNX Runtime C++端侧模型部署YOLOv5
- Ai
- 3天前
- 30热度
- 0评论
加载准备
初始化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