OnnxRunTime的部署流程
文章目录
-
- 一、torch转onnx模型
- 二、python下推断onnx模型
- 三、C++下用OnnxRunTime来推断模型
一、torch转onnx模型
用pytorch训练得到模型后，这时候需要转成onnx模型
torch提供函数
torch.onnx.export()
def export():
# 自己定义的网络结构
model = CSRNet_onnx(load_weights=r'./weights/mutil_hjxu_d_4.78.path.tar')
#
pthfile = r'./weights/mutil_hjxu_d_4.78.path.tar'
checkpoint = torch.load(pthfile, map_location="cpu")
model.load_state_dict(checkpoint['state_dict'])
# data type nhwc 1152, 648
dummy_input1 = torch.randn(1, 540, 960, 3)
input_names = ["input_1"]
output_names = ["output1"]
dynamic_axes = {'input_1': [0, 1, 2]}
torch.onnx.export(model, dummy_input1, "./weights/csrNet_540*960_hjxu.onnx", verbose=True, input_names=input_names,
output_names=output_names)
# 动态大小
# torch.onnx.export(model, dummy_input1, "./weights/model540*960_no.onnx", verbose=True, input_names=input_names,
# output_names=output_names,opset_version=10, dynamic_axes=dynamic_axes)
二、python下推断onnx模型
从python段看，onnx推断模型，分以下几个步骤
- step1：加载模型
- step2：获取输入输出节点名
- step3：准备数据，如（n，c，h，w）或者（n， h， w， c）
- step4：执行推断，获得输出
import torch
import onnxruntime as rt
import torch.onnx
from model import CSRNet, CSRNet_onnx
import cv2
import numpy as np
import onnx
def test_onnx():
img_paths = ["./result/000000000192.jpg"]
sess = rt.InferenceSession("./weights/csrNet_540*960_hjxu.onnx",None)
input_name = sess.get_inputs()[0].name
out_name = sess.get_outputs()[0].name
# 打印输入节点的名字，以及输入节点的shape
for i in range(len(sess.get_inputs())):
print(sess.get_inputs()[i].name, sess.get_inputs()[i].shape)
print("----------------")
# 打印输出节点的名字，以及输出节点的shape
for i in range(len(sess.get_outputs())):
print(sess.get_outputs()[i].name, sess.get_outputs()[i].shape)
for img_path in img_paths:
X_test = cv2.imread(img_path)
X_test = cv2.resize(X_test, (960, 540))
X_test = cv2.cvtColor(X_test, cv2.COLOR_BGR2RGB)
x = []
x.append(X_test)
pred_onx = sess.run([out_name], {input_name:x}) # 执行推断
三、C++下用OnnxRunTime来推断模型
这里笔者使用c的接口来推断
先介绍一下大致流程：
- step1、首先得获取一个OrtApi对象。这个主要是由于onnxruntime在c api的接口实现中，将接口都初始化一个全局静态OrtApi对象，所以这些函数接口指针并不需要一个一个去获取，只用获取到这个OrtApi对象即可通过访问成员的方式使用对应的函数，就可以了，这极大极大方便了接口使用。而获取这个对象是通过一个OrtGetApiBase函数实现。
const OrtApi* g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION);
- step2、初始化一个环境变量，可以用来创建onnxruntime的session
OrtEnv* env;
- step3、创建一个tensor，这里不得不提到OrtValue。OrtValue 是onnxruntime c api中最常见的类型，用来封装所有的数据传递结构。其中最主要的就是tensor，所以api中也提供了一系列相关的函数接口用于tensor相关的类型转换。
由于我们是在ONNX中调用onnxruntime的api，而且是用动态加载的方式，所以我们只用关心从onnx tensor 到OrtValue的相互转换，而OrtValue内部封装的是onnxruntime的tensor类型。
CreateTensorAsOrtValue / CreateTensorWithDataAsOrtValue 这两个函数进行OrtValue类型的tensor创建，这两个函数的区别一是是否需要由onnxruntime进行内存分配及其内存管理的职责。
-step4、推断。就和python中的sess.run()就可以了
g_ort->Run(session, NULL, input_node_names.data(), (const OrtValue* const*)&input_tensor, 1, output_node_names.data(), 1, &output_tensor)
整体推断代码官网demo就可见，这里简单复制一下。
#include >
printf("Number of inputs = %zu\n", num_input_nodes);
// iterate over all input nodes
for (size_t i = 0; i < num_input_nodes; i++) {
// print input node names
char* input_name;
status = g_ort->SessionGetInputName(session, i, allocator, &input_name);
printf("Input %zu : name=%s\n", i, input_name);
input_node_names[i] = input_name;
// print input node types
OrtTypeInfo* typeinfo;
status = g_ort->SessionGetInputTypeInfo(session, i, &typeinfo);
const OrtTensorTypeAndShapeInfo* tensor_info;
CheckStatus(g_ort->CastTypeInfoToTensorInfo(typeinfo, &tensor_info));
ONNXTensorElementDataType type;
CheckStatus(g_ort->GetTensorElementType(tensor_info, &type));
printf("Input %zu : type=%d\n", i, type);
// print input shapes/dims
size_t num_dims;
CheckStatus(g_ort->GetDimensionsCount(tensor_info, &num_dims));
printf("Input %zu : num_dims=%zu\n", i, num_dims);
input_node_dims.resize(num_dims);
g_ort->GetDimensions(tensor_info, (int64_t*)input_node_dims.data(), num_dims);
for (size_t j = 0; j < num_dims; j++)
printf("Input %zu : dim %zu=%jd\n", i, j, input_node_dims[j]);
g_ort->ReleaseTypeInfo(typeinfo);
}
// Results should be...
// Number of inputs = 1
// Input 0 : name = data_0
// Input 0 : type = 1
// Input 0 : num_dims = 4
// Input 0 : dim 0 = 1
// Input 0 : dim 1 = 3
// Input 0 : dim 2 = 224
// Input 0 : dim 3 = 224
//*************************************************************************
// Similar operations to get output node information.
// Use OrtSessionGetOutputCount(), OrtSessionGetOutputName()
// OrtSessionGetOutputTypeInfo() as shown above.
//*************************************************************************
// Score the model using sample data, and inspect values
size_t input_tensor_size = 224 * 224 * 3; // simplify ... using known dim values to calculate size
// use OrtGetTensorShapeElementCount() to get official size!
std::vector<float> input_tensor_values(input_tensor_size);
std::vector<const char*> output_node_names = { "output1" }; // 输出节点
// initialize input data with values in [0.0, 1.0]
// 这里用的是直接赋值
// input_tensor_values.data() 相当于 mat.data()
for (size_t i = 0; i < input_tensor_size; i++)
input_tensor_values[i] = (float)i / (input_tensor_size + 1);
// create input tensor object from data values
OrtMemoryInfo* memory_info;
CheckStatus(g_ort->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &memory_info));
OrtValue* input_tensor = NULL;
g_ort->CreateTensorAsOrtValue(allocator, input_node_dims.data(), 4, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &input_tensor);
CheckStatus(g_ort->CreateTensorWithDataAsOrtValue(memory_info, input_tensor_values.data(), input_tensor_size * sizeof(float), input_node_dims.data(), 4, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &input_tensor));
int is_tensor;
CheckStatus(g_ort->IsTensor(input_tensor, &is_tensor));
assert(is_tensor);
g_ort->ReleaseMemoryInfo(memory_info);
// score model & input tensor, get back output tensor
OrtValue* output_tensor = NULL;
//while (true)
//{
CheckStatus(g_ort->Run(session, NULL, input_node_names.data(), (const OrtValue* const*)&input_tensor, 1, output_node_names.data(), 1, &output_tensor));
CheckStatus(g_ort->IsTensor(output_tensor, &is_tensor));
assert(is_tensor);
// Get pointer to output tensor float values
float* floatarr;
CheckStatus(g_ort->GetTensorMutableData(output_tensor, (void**)&floatarr));
assert(std::abs(floatarr[0] - 0.000045) < 1e-6);
// score the model, and print scores for first 5 classes
for (int i = 0; i < 5; i++)
printf("Score for class [%d] = %f\n", i, floatarr[i]);
//}
g_ort->ReleaseValue(output_tensor);
g_ort->ReleaseValue(input_tensor);
g_ort->ReleaseSession(session);
g_ort->ReleaseSessionOptions(session_options);
g_ort->ReleaseEnv(env);
printf("Done!\n");
return 0;
}
码字不易，看完如有帮助，点个关注呗~