利用OnnxRuntime进行torch模型部署(C++版)——以分类网络为例
近年来,ONNX(Open Neural Network Exchange)格式的兴起为解决这一问题提供了有效的途径。ONNX是一种开放标准,旨在使不同深度学习框架之间能够共享模型。通过将模型导出为ONNX格式,我们可以轻松地在不同框架和平台上进行部署和推理。而ONNX Runtime,作为微软开源的高性能推理引擎,更是为模型的部署提供了强大的支持。
本文将重点介绍如何利用ONNX Runtime在C++环境中部署PyTorch模型。PyTorch以其灵活性和易用性在研究和开发领域广受欢迎,而C++则因其高性能和稳定性在生产环境中占据重要位置。通过结合这两者的优势,我们可以实现深度学习模型的高效部署和推理。
windows安装onnxruntime
安装地址:
windows安装onnxruntime
然后在tags中去找需要下载的版本,我这里下载的是v1.15.0 。找到windows GPU版推理版,如果你的是CPU的就下载CPU即可。
注意:在后面配置环境时,需要在Release下配置。
下载完成并解压后,如下图所示:
Visual Studio配置OnnxRuntime
我这里是Visual Studio 2017。这里新建项目,我这里项目名称为“Project1”.然后在项目源文件中新建一个main.cpp文件。
配置include路径
进入项目属性,在Release|x64-->C/C++-->常规-->附件包含目录添加之前onnxrutime-gpu的include路径,如下:
配置链接
配置链接(静态库),点击链接器-->输入-->附加依赖项,添加onnxruntime的lib文件路径,如下:
我这里添加的为(包含了Opencv):
onnxruntime.lib
onnxruntime_providers_cuda.lib
onnxruntime_providers_shared.lib
opencv_world410.lib
复制动态库到项目下
完成上述的配置后,将onnxruntime中的onnxruntime.dll、onnxruntime_providers_cuda.dll和onnxruntime_providers_shared.dll以及Opencv的dll复制到自己的Release下。(如果你没有x64/Release文件夹,你可以随便写的主函数后run一下就有了,注意是在Release下,不是Debug)
完成上述的操作后,我们进可用在代码中加入onnx的头文件了,如下:
pytorch2onnx模型
再利用c++调用训练好的模型前,我们需要用pytorch先导出onnx模型,torch官方给我们提供了非常遍历的api,代码示例如下,我这里仅以一个非常简单的猫狗分类网络为例:
我的网络为mobilnetv2,然后利用torch.onnx.export导出onnx即可,这里需要注意输入输出节点名字,因为我这里是只有一个输入、一个输出,输入节点名字为['image0'],输出为['output'],注意这里要用[ ]引用起来,即便你只有一个输入输出。
import cv2
import torch
from nets import mobilenet_v2
if __name__ =='__main__':
x = torch.zeros(1, 3, 224, 224).cuda()
model = mobilenet_v2(pretrained=False, num_classes=2)
ckpt_path = 'model_data/mobilenet_catvsdog.pth'
print(model.load_state_dict(torch.load(ckpt_path,map_location='cpu')))
model = model.cuda()
model.eval()
torch.onnx.export(model, x, 'mobilenet_v2.onnx',verbose=True,input_names=['image0'],output_names=['output'])
C++调用onnx模型
模型的加载
完成了torch转onnx的操作,接下来就可以用C++调用了,代码如下:
session_options:注意是做一些设置。
env是创建的onnxruntime运行环境。
Ort::session session_是创建的一个会话,该会话可以执行已经加载到该会话中的onnx模型。具体可以完成:
-
接受输入数据:会话可以接收用户提供的输入数据,这些数据通常是以张量(Tensor)的形式提供的,用于模型的推理计算。
-
执行模型推理:在接收到输入数据后,会话会根据加载的ONNX模型执行推理计算,包括前向传播等必要的数学运算,以生成输出结果。
-
获取输出结果:推理计算完成后,会话会生成输出张量,这些输出张量包含了模型的推理结果,用户可以从会话中获取这些结果并进行后续处理。
-
管理会话状态:会话还负责管理模型加载后的状态,包括模型的输入输出接口、配置参数等,以确保模型能够正确、高效地执行推理任务。
std::string onnxpath = "E:/classification-pytorch-main/mobilenet_v2.onnx"; //onnx路径
std::wstring modelPath = std::wstring(onnxpath.begin(), onnxpath.end());
Ort::SessionOptions session_options;
Ort::Env env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "model_load_err");
// 设定单个操作(op)内部并行执行的最大线程数,可以提升速度
session_options.SetIntraOpNumThreads(20);
session_options.SetGraphOptimizationLevel(ORT_ENABLE_EXTENDED);
std::cout << "onnxruntime inference try to use GPU Device" << std::endl;
// 是否使用GPU,以及GPU ID
int device_id = 0;
std::cout << "当前GPU ID 为" << device_id << std::endl;
OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, device_id); //可选择GPU ID
Ort::Session session_(env, modelPath.c_str(), session_options);
获得输入输出节点
利用Ort::Session创建的session_对象完成了onnx模型的加载,我们就可以进一步利用该会话去获得我们的输入输出节点了:
GetInputCount()和GetOutputCount()可以用来获取输入输出的节点数量。
int input_nodes_num = session_.GetInputCount(); //获取输入数量
int output_nodes_num = session_.GetOutputCount(); //获取输出数量
要进一步获取输入输出节点的名字,我们还需要创建一个分配器,该分配器可以分配和释放内存,管理着模型推理过程中需要的内存资源:
Ort::AllocatorWithDefaultOptions allocator; //分配器,分配和和释放内存,管理着模型推理过程所需要的内存资源
然后通过遍历输入节点,配合着上述的分配器,就可以获取输入输出节点的名字。这里只放输入节点的获取,输出的是一样的。
利用GetInputNameAllocated()可以获取输入节点,然后使用get()方法获取输入节点名字的字符串。利用GetInputTypeInfor获取输入信息,利用GetTensorTypeAndShape().GetShape()可以获取张量的形状。
std::vector<std::string> input_node_names; //获取输入的节点name
std::vector<std::string> output_node_names; //获取输出的节点name
//遍历每个输入节点
for (int i = 0; i < input_nodes_num; i++)
{
auto input_name = session_.GetInputNameAllocated(i, allocator);
input_node_names.push_back(input_name.get());
auto inputShapeInfo = session_.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape(); //输入的shape (channels, high, width)
int ch = inputShapeInfo[1];
input_h = inputShapeInfo[2];
input_w = inputShapeInfo[3];
//std::cout << "input format: " << ch << "x" << input_h << "x" << input_w << std::endl;
}
如果你希望打印输入输出节点的名字,以和你pytorch进行对应,可以进行如下操作:
std::cout << "输入节点:" << input_node_names[0].c_str() << std::endl;
std::cout << "输出节点:" << output_node_names[0].c_str() << std::endl;
推理
然后我们就可以进行图像的推理了(注意,这里的图像是预处理好的,我会在后面章节给出图像处理的代码,这里只展示onnx推理的核心代码)。
allocator_info:是在CPU上创建一块运行区域,也就是将最后的结果从GPU搬到CPU。
input_tensor_:是我们创建的张量,也就是把图像转为tensor,适合模型的输入。
然后调用session_会话的Run方法,进行推理,主要传入:输入节点的地址(inputNames.data()),告诉模型从这里输入,输入的东西是什么呢?输入的是前面创建的input_tensro_(地址引用),然后给出输出的地址和大小,也就是告诉模型从哪里进来,从哪里出去。
ort_outputs就是我们获得输出。
auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); //在cpu上创建运行区域
Ort::Value input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, blob.ptr<float>(), tpixels, input_shape_info.data(), input_shape_info.size());
std::vector<Ort::Value> ort_outputs;
ort_outputs = session_.Run(Ort::RunOptions{ nullptr }, inputNames.data(), &input_tensor_, inputNames.size(), outNames.data(), outNames.size());
那么如何进一步处理输出呢?比如我这里是猫狗分类,两个类,我要获取,也就是我torch中的输出shape为(batch_size,num_cls),也就是(1,2)。
获取输出数据信息(地址)
利用以下代码就可以获取输出的数据,也就是调用GetTensorData。这里的ort_outputs[0]是获取的第一个输出(注意不是第一个分类,是该输出包含了两个类),*output_data是指向该输出第一个类的地址。
const float* output_data = ort_outputs[0].GetTensorData<float>();
获取输出的shape
用GetTensorTypeAndShapeInfo().GetShape().data()可以获得输出张量的shape,
int64_t* output_shape = ort_outputs[0].GetTensorTypeAndShapeInfo().GetShape().data();
int batch_size = output_shape[0];
int num_classes = output_shape[1];
std::cout << "num_classes: "<< num_classes << std::endl;
softmax处理
由于我们的输出没有进行softmax处理,因此需要进行如下处理:
// 将输出转换为 cv::Mat
cv::Mat logits(batch_size, num_classes, CV_32F, (void*)output_data);
// 应用 softmax
cv::Mat probabilities = softmax(logits);
softmax函数的定义如下:
// Softmax 函数
cv::Mat softmax(const cv::Mat& logits) {
CV_Assert(logits.type() == CV_32F);
int batch_size = logits.size[0];
int num_classes = logits.size[1];
cv::Mat result(batch_size, num_classes, CV_32F);
for (int i = 0; i < batch_size; ++i) {
const float* logit_ptr = logits.ptr<float>(i);
float* result_ptr = result.ptr<float>(i);
// 计算每个样本的 softmax
float max_val = *(std::max_element(logit_ptr, logit_ptr + num_classes));
float sum_exp = 0.0f;
for (int j = 0; j < num_classes; ++j) {
result_ptr[j] = std::exp(logit_ptr[j] - max_val);
sum_exp += result_ptr[j];
}
for (int j = 0; j < num_classes; ++j) {
result_ptr[j] /= sum_exp;
}
}
return result;
}
接着就可以输出最大的预测概率:
// 输出最大概率的类别
for (int i = 0; i < batch_size; ++i) {
int max_idx = std::max_element(probabilities.ptr<float>(i), probabilities.ptr<float>(i) + num_classes) - probabilities.ptr<float>(i);
std::cout << "Sample " << i << " predicted class: " << cls_names[max_idx] << " with probability " << probabilities.at<float>(i, max_idx) << std::endl;
}
输出如下:
onnxruntime inference try to use GPU Device
当前GPU ID 为0
output format: 2x3x-331813314
preprocessedImage shape:(224,224,3)
(1,3,224,224)
输入节点:image0
输出节点:output
num_classes: 2
[-3.6467652, 4.0626707]
Sample 0 predicted class: dog with probability 0.999552
而我们输入的图像正好是dog,如下:
图像的预处理
现在来说一下在c++中如何处理我们的图像,使其适合我们onnx的输入。
图像不失真的reshape
图像不失真的reshape如下,这个仿照python中写的:
cv::Mat letterbox_image_opencv(cv::Mat image, const cv::Size& size)
{
//BGR2RGB
cv::Mat rgb_image;
cv::cvtColor(image, rgb_image, cv::COLOR_BGR2RGB);
//获取图像的原始大小
int ih = image.rows;
int iw = image.cols;
int h = size.height;
int w = size.width;
double scale = std::min(static_cast<double>(w) / iw, static_cast<double>(h) / ih);
// 计算调整大小后的图像宽度和高度
int nw = static_cast<int>(iw * scale);
int nh = static_cast<int>(ih * scale);
cv::Mat resizedImage;
//reshape目标大小
cv::resize(image, resizedImage, cv::Size(nw,nh),0,0,cv::INTER_CUBIC);
//创建新的图像,填充为灰色(128,128,128)
cv::Mat newImage = cv::Mat::zeros(size, CV_8UC3) + cv::Scalar(128, 128, 128);
// 计算填充的起始位置,使调整大小后的图像居中
int startX = (w - nw) / 2;
int startY = (h - nh) / 2;
// 将调整大小后的图像复制到新图像的中央位置
resizedImage.copyTo(newImage(cv::Rect(startX, startY, nw, nh)));
return newImage;
}
图像归一化
cv::Mat preprocess_input(const cv::Mat& image)
{
cv::Mat preprocessed_image;
image.convertTo(preprocessed_image, CV_32F);
// 归一化像素值到[-1, 1]范围
preprocessed_image = preprocessed_image / 127.5f - 1.0f;
return preprocessed_image;
}
HWC转NCHW
这里在添加batch维度和进行通道的转换时,用的是opencv中的cv::dnn::blobFromImage方法。
如下:
cv::Size size(224, 224);
cv::Mat image_data = letterbox_image_opencv(img, size);//reshape image
cv::Mat preprocessedImage = preprocess_input(image_data); //归一化后的图像,shape 为(224,224, 3)
cv::Mat blob = cv::dnn::blobFromImage(preprocessedImage, 1.0, size, cv::Scalar(0, 0, 0), false, false); //blob为4D矩阵,没有定义行/列的值
则我们最终创建的输入为:
这样我们就得到最终的输入shape为(1,3,224,224)
cv::Mat blob = cv::dnn::blobFromImage(preprocessedImage, 1.0, size, cv::Scalar(0, 0, 0), false, false); //blob为4D矩阵,没有定义行/列的值
//获取blob的shape
int bs = blob.size[0];
int channels = blob.size[1];
int height = blob.size[2];
int width = blob.size[3];
std::cout << "(" << bs << "," << channels << "," << height << "," << width << ")" << std::endl;
size_t tpixels = size.width * size.height * 3; //224*224*3
std::array<int64_t, 4> input_shape_info{ bs, 3, size.height, size.width };
完整代码
#include <stdio.h>
#include <string.h>
#include <opencv2/opencv.hpp>
#include <fstream>
#include "onnxruntime_cxx_api.h"
const std::vector<std::string> cls_names = { "cat","dog" };
cv::Mat letterbox_image_opencv(cv::Mat image, const cv::Size& size)
{
//BGR2RGB
cv::Mat rgb_image;
cv::cvtColor(image, rgb_image, cv::COLOR_BGR2RGB);
//获取图像的原始大小
int ih = image.rows;
int iw = image.cols;
int h = size.height;
int w = size.width;
double scale = std::min(static_cast<double>(w) / iw, static_cast<double>(h) / ih);
// 计算调整大小后的图像宽度和高度
int nw = static_cast<int>(iw * scale);
int nh = static_cast<int>(ih * scale);
cv::Mat resizedImage;
//reshape目标大小
cv::resize(image, resizedImage, cv::Size(nw,nh),0,0,cv::INTER_CUBIC);
//创建新的图像,填充为灰色(128,128,128)
cv::Mat newImage = cv::Mat::zeros(size, CV_8UC3) + cv::Scalar(128, 128, 128);
// 计算填充的起始位置,使调整大小后的图像居中
int startX = (w - nw) / 2;
int startY = (h - nh) / 2;
// 将调整大小后的图像复制到新图像的中央位置
resizedImage.copyTo(newImage(cv::Rect(startX, startY, nw, nh)));
return newImage;
}
cv::Mat preprocess_input(const cv::Mat& image)
{
cv::Mat preprocessed_image;
image.convertTo(preprocessed_image, CV_32F);
// 归一化像素值到[-1, 1]范围
preprocessed_image = preprocessed_image / 127.5f - 1.0f;
return preprocessed_image;
}
// Softmax 函数
cv::Mat softmax(const cv::Mat& logits) {
CV_Assert(logits.type() == CV_32F);
int batch_size = logits.size[0];
int num_classes = logits.size[1];
cv::Mat result(batch_size, num_classes, CV_32F);
for (int i = 0; i < batch_size; ++i) {
const float* logit_ptr = logits.ptr<float>(i);
float* result_ptr = result.ptr<float>(i);
// 计算每个样本的 softmax
float max_val = *(std::max_element(logit_ptr, logit_ptr + num_classes));
float sum_exp = 0.0f;
for (int j = 0; j < num_classes; ++j) {
result_ptr[j] = std::exp(logit_ptr[j] - max_val);
sum_exp += result_ptr[j];
}
for (int j = 0; j < num_classes; ++j) {
result_ptr[j] /= sum_exp;
}
}
return result;
}
int main()
{
cv::Mat img = cv::imread("E:/classification-pytorch-main/img/dog.jpg"); //读取图像
if (img.empty())
{
printf("读取图像错误");
return -1;
}
std::string onnxpath = "E:/classification-pytorch-main/mobilenet_v2.onnx"; //onnx路径
std::wstring modelPath = std::wstring(onnxpath.begin(), onnxpath.end());
Ort::SessionOptions session_options;
Ort::Env env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "model_load_err");
// 设定单个操作(op)内部并行执行的最大线程数,可以提升速度
session_options.SetIntraOpNumThreads(20);
session_options.SetGraphOptimizationLevel(ORT_ENABLE_EXTENDED);
std::cout << "onnxruntime inference try to use GPU Device" << std::endl;
// 是否使用GPU,以及GPU ID
int device_id = 0;
std::cout << "当前GPU ID 为" << device_id << std::endl;
OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, device_id); //可选择GPU ID
Ort::Session session_(env, modelPath.c_str(), session_options);
int input_nodes_num = session_.GetInputCount(); //获取输入数量
int output_nodes_num = session_.GetOutputCount(); //获取输出数量
std::vector<std::string> input_node_names; //获取输入的节点name
std::vector<std::string> output_node_names; //获取输出的节点name
Ort::AllocatorWithDefaultOptions allocator; //分配器,分配和和释放内存,管理着模型推理过程所需要的内存资源
int input_h = 0;
int input_w = 0;
int output_h = 0;
int output_w = 0;
//遍历每个输入节点
for (int i = 0; i < input_nodes_num; i++)
{
auto input_name = session_.GetInputNameAllocated(i, allocator);
input_node_names.push_back(input_name.get());
auto inputShapeInfo = session_.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape(); //输入的shape (channels, high, width)
int ch = inputShapeInfo[1];
input_h = inputShapeInfo[2];
input_w = inputShapeInfo[3];
//std::cout << "input format: " << ch << "x" << input_h << "x" << input_w << std::endl;
}
//遍历每个输出节点
for (int i = 0; i < output_nodes_num; i++)
{
auto output_name = session_.GetOutputNameAllocated(i, allocator);
output_node_names.push_back(output_name.get());
auto outputShapeInfo = session_.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape(); //输出的shape
int ch_out = outputShapeInfo[1];
output_h = outputShapeInfo[2];
output_w = outputShapeInfo[3];
std::cout << "output format: " << ch_out << "x" << output_h << "x" << output_w << std::endl;
}
cv::Size size(224, 224);
cv::Mat image_data = letterbox_image_opencv(img, size);//reshape image
cv::Mat preprocessedImage = preprocess_input(image_data); //归一化后的图像,shape 为(224,224, 3)
std::cout << "preprocessedImage shape:" << "("<< preprocessedImage.size[0] << ","<< preprocessedImage.size[1] <<"," << preprocessedImage.channels() << ")"<<std::endl;
cv::Mat blob = cv::dnn::blobFromImage(preprocessedImage, 1.0, size, cv::Scalar(0, 0, 0), false, false); //blob为4D矩阵,没有定义行/列的值
//获取blob的shape
int bs = blob.size[0];
int channels = blob.size[1];
int height = blob.size[2];
int width = blob.size[3];
std::cout << "(" << bs << "," << channels << "," << height << "," << width << ")" << std::endl;
size_t tpixels = size.width * size.height * 3; //224*224*3
std::array<int64_t, 4> input_shape_info{ bs, 3, size.height, size.width };
auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); //在cpu上创建运行区域
Ort::Value input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, blob.ptr<float>(), tpixels, input_shape_info.data(), input_shape_info.size());
std::cout << "输入节点:" << input_node_names[0].c_str() << std::endl;
std::cout << "输出节点:" << output_node_names[0].c_str() << std::endl;
// 输入一个数据
const std::array<const char*, 1> inputNames = { input_node_names[0].c_str() };
// 输出一个数据
const std::array<const char*, 1> outNames = { output_node_names[0].c_str()};
std::vector<Ort::Value> ort_outputs;
try
{
ort_outputs = session_.Run(Ort::RunOptions{ nullptr }, inputNames.data(), &input_tensor_, inputNames.size(), outNames.data(), outNames.size());
// 假设输出是一个包含 logits 的 tensor
const float* output_data = ort_outputs[0].GetTensorData<float>();//只有一个输出,所以是ort_outputs[0],不是指输出几个类的第一个类
//std::cout << "*output_data:" << *output_data << std::endl;
int64_t* output_shape = ort_outputs[0].GetTensorTypeAndShapeInfo().GetShape().data();
int batch_size = output_shape[0];
int num_classes = output_shape[1];
std::cout << "num_classes: "<< num_classes << std::endl;
// 将输出转换为 cv::Mat
cv::Mat logits(batch_size, num_classes, CV_32F, (void*)output_data);
std::cout << logits << std::endl;
// 应用 softmax
cv::Mat probabilities = softmax(logits);
// 输出最大概率的类别
for (int i = 0; i < batch_size; ++i) {
int max_idx = std::max_element(probabilities.ptr<float>(i), probabilities.ptr<float>(i) + num_classes) - probabilities.ptr<float>(i);
std::cout << "Sample " << i << " predicted class: " << cls_names[max_idx] << " with probability " << probabilities.at<float>(i, max_idx) << std::endl;
}
}
catch (const std::exception& e)
{
std::cout << e.what() << std::endl;
}
return 0;
}