Pytorch导出onnx模型并在C++环境中调用（含python和C++工程）

工程下载链接：Pytorch导出onnx模型并在C++环境中调用（python和C++工程）

机器学习多层感知机MLP的Pytorch实现-以表格数据为例-含数据集和PyCharm工程中简单介绍了在python中使用pytorch搭建神经网络模型的步骤和代码工程，此处介绍AI模型的跨平台调用问题，即使用跨平台的ONNX框架，在C++代码中进行模型调用。

参考：pytorch导出模型并使用onnxruntime C++部署加载模型推理

1、pkl权重文件转换为onnx模型

在机器学习多层感知机MLP的Pytorch实现-以表格数据为例-含数据集和PyCharm工程中，我们对训练完成的模型进行了模型的保存：

torch.save(model.state_dict(),'weights/mlp_weights-epoch%d-Total_loss%.4f-val_loss%.4f.pkl' % ((epoch + 1), train_loss, val_loss / (iteration + 1)))

此处我们需要先加载保存的模型，如何再将其导出为onnx格式。

1.1、加载保存的模型

这一步主要是要把保存的模型恢复出来：

import numpy as np
import onnxruntime
import torch
from torch import nn
import torch.nn.functional as F# 定义多层感知机（MLP）类，继承自nn.Module
class MLP(nn.Module):# 类的初始化方法def __init__(self):# 调用父类nn.Module的初始化方法super(MLP, self).__init__()self.hidden1 = nn.Linear(in_features=8, out_features=50, bias=True)self.hidden2 = nn.Linear(50, 50)self.hidden3 = nn.Linear(50, 50)self.hidden4 = nn.Linear(50, 50)self.predict = nn.Linear(50, 1)# 定义前向传播方法def forward(self, x):# x是输入数据，通过第一个隐藏层并应用ReLU激活函数x = F.relu(self.hidden1(x))x = F.relu(self.hidden2(x))x = F.relu(self.hidden3(x))x = F.relu(self.hidden4(x))output = self.predict(x)# 将输出展平为一维张量# out = output.view(-1)return output# 检查是否有可用的GPU
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = "cpu"
# 定义模型并将其转移到GPU
model = MLP().to(device)
model.load_state_dict(torch.load('weights/mlp_weights-epoch1000-Total_loss0.0756-val_loss0.0105.pkl',weights_only=True))
model.eval()

1.2、pytorch调用加载的模型进行测试

先简单测试一下这个模型，定义一个输入全为0的数组作为输入，打印输出的结果：

# 初始化一个输入全为0的数组进行测试
x = (torch.from_numpy(np.array([0, 0, 0, 0, 0, 0, 0, 0]).astype(np.float32)).to(device))
y = model(x).cpu().detach().numpy()[0]
print(f"pytorch直接测试结果为: {y}")

1.3、导出onnx模型

使用下面的代码将原来的模型model导出为onnx的模型，其中x是上面定义的案例输入：

batch_size = 1  # 批处理大小
export_onnx_file = "test.onnx"  # 目的ONNX文件名
torch.onnx.export(model,(x),export_onnx_file,opset_version=10,do_constant_folding=True,  # 是否执行常量折叠优化input_names=["input"],  # 输入名output_names=["output"],  # 输出名dynamic_axes={"input": {0: "batch_size"},  # 批处理变量"output": {0: "batch_size"}})

这个函数的具体定义可以参考：从pytorch转换到onnx

1.4、在python中调用onnx模型测试

使用下面代码加载onnx模型并进行测试：

resnet_session = onnxruntime.InferenceSession(export_onnx_file)
inputs = {resnet_session.get_inputs()[0].name: x.cpu().detach().numpy()}
outs = resnet_session.run(None, inputs)[0][0]
print(f"py onnx直接测试结果为: {outs}")

可以看到pytorch的结果和onnx的结果是基本一致的：

1.5、全部代码

import numpy as np
import onnxruntime
import torch
from torch import nn
import torch.nn.functional as F# 定义多层感知机（MLP）类，继承自nn.Module
class MLP(nn.Module):# 类的初始化方法def __init__(self):# 调用父类nn.Module的初始化方法super(MLP, self).__init__()self.hidden1 = nn.Linear(in_features=8, out_features=50, bias=True)self.hidden2 = nn.Linear(50, 50)self.hidden3 = nn.Linear(50, 50)self.hidden4 = nn.Linear(50, 50)self.predict = nn.Linear(50, 1)# 定义前向传播方法def forward(self, x):# x是输入数据，通过第一个隐藏层并应用ReLU激活函数x = F.relu(self.hidden1(x))x = F.relu(self.hidden2(x))x = F.relu(self.hidden3(x))x = F.relu(self.hidden4(x))output = self.predict(x)# 将输出展平为一维张量# out = output.view(-1)return output# 检查是否有可用的GPU
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = "cpu"
# 定义模型并将其转移到GPU
model = MLP().to(device)
model.load_state_dict(torch.load('weights/mlp_weights-epoch1000-Total_loss0.0756-val_loss0.0105.pkl',weights_only=True))
model.eval()# 初始化一个输入全为0的数组进行测试
x = (torch.from_numpy(np.array([0, 0, 0, 0, 0, 0, 0, 0]).astype(np.float32)).to(device))
y = model(x).cpu().detach().numpy()[0]
print(f"pytorch直接测试结果为: {y}")batch_size = 1  # 批处理大小
export_onnx_file = "test.onnx"  # 目的ONNX文件名
torch.onnx.export(model,(x),export_onnx_file,opset_version=10,do_constant_folding=True,  # 是否执行常量折叠优化input_names=["input"],  # 输入名output_names=["output"],  # 输出名dynamic_axes={"input": {0: "batch_size"},  # 批处理变量"output": {0: "batch_size"}})resnet_session = onnxruntime.InferenceSession(export_onnx_file)
inputs = {resnet_session.get_inputs()[0].name: x.cpu().detach().numpy()}
outs = resnet_session.run(None, inputs)[0][0]
print(f"py onnx直接测试结果为: {outs}")

2、C++调用onnx模型

2.1、库的下载安装和官方手册

这个整个库的下载还是要到官方的github仓库去：microsoft/onnxruntime
具体的使用方式参考英文的手册：https://onnxruntime.ai/docs/】

此处下载的window版本的，下载下来可以得到头文件和库文件：

因此在实际编程的时候我使用的Cmakelist来链接到相关的库，我是使用VS code + gcc构成的C++编译环境：

if (WIN32)include_directories(${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-1.14.0/onnxruntime-win-x64-1.14.0/include)
else()include_directories(${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-1.14.0/onnxruntime-linux-x64-1.14.0/include)
endif()if (WIN32)link_directories(${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-1.14.0/onnxruntime-win-x64-1.14.0/lib)
else()link_directories(${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-1.14.0/onnxruntime-linux-x64-1.14.0/lib)
endif()

实际的工程目录结构如下所示：

2.2、C++调用代码实现

下面代码实现和1.4、在python中调用onnx模型测试相同的效果，输入是全0的数组，进行计算并返回相关结果：

#include <iostream>
#include <array>
#include <algorithm>
#include "onnxruntime_cxx_api.h"#define ONNX_IN_OUT_SIZE_MAX 20int main(int argc, char* argv[])
{//print helloprintf("hello");std::vector<float> input_matrix_vector={0, 0, 0, 0, 0, 0, 0, 0};int onnx_input_shape = 8;int onnx_output_shape = 1;// --- define model path#if _WIN32const wchar_t* model_path = L"./model.onnx"; // you can use string to wchar_t* function to convert#elseconst char* model_path = "./model.onnx";#endif// --- init onnxruntime envOrt::Env env(ORT_LOGGING_LEVEL_WARNING, "Default");auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);// set optionsOrt::SessionOptions session_option;session_option.SetIntraOpNumThreads(1); // extend the number to do parallelsession_option.SetGraphOptimizationLevel(ORT_ENABLE_ALL);// --- prepare dataconst char* input_names[] = { "input" }; // must keep the same as model exportconst char* output_names[] = { "output" };std::array<float, ONNX_IN_OUT_SIZE_MAX> input_matrix;std::array<float, ONNX_IN_OUT_SIZE_MAX> output_matrix;if(input_matrix_vector.size()>ONNX_IN_OUT_SIZE_MAX){throw std::runtime_error("input_matrix_vector.size()<ONNX_IN_OUT_SIZE_MAX");}std::copy(input_matrix_vector.begin(),input_matrix_vector.end(),input_matrix.begin());// must use int64_t type to match argsstd::array<int64_t, 1> input_shape{ onnx_input_shape };std::array<int64_t, 1> output_shape{ onnx_output_shape };Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_matrix.data(), input_matrix.size(), input_shape.data(), input_shape.size());Ort::Value output_tensor = Ort::Value::CreateTensor<float>(memory_info, output_matrix.data(), output_matrix.size(), output_shape.data(), output_shape.size());// --- predictOrt::Session session(env, model_path, session_option); // FIXME: must check if model file exist or valid, otherwise this will cause crashsession.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor, 1, output_names, &output_tensor, 1); // here only use one input output channelstd::vector<float> outputVector(onnx_output_shape);std::copy(output_matrix.begin(),output_matrix.begin()+onnx_output_shape,outputVector.begin());std::cout << "--- predict result ---" << std::endl;// matrix outputstd::cout << "ouput matrix: ";for (int i = 0; i < outputVector.size(); i++)std::cout << outputVector[i] << " ";std::cout << std::endl;// argmax value// int argmax_value = std::distance(output_matrix.begin(), std::max_element(output_matrix.begin(), output_matrix.end()));// std::cout << "output argmax value: " << argmax_value << std::endl;// getchar();return 0;
}

可以看到最终的返回结果为：

和之前在python中的结果是一致的！！！

2.3、注意，模型文件和onnx的dll要在exe同一级目录

3、运行时遇到的一些问题

编译报错—error: ‘Frees_ptr_opt‘ has not been declared

在编译器命令行或者代码中定义这些宏，使其在非MSVC环境中被忽略。在代码的开头（ onnxruntime_c_api.h 文件中）添加以下代码：

#pragma once
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
//add code here
#ifndef _Frees_ptr_opt_ 
#define _Frees_ptr_opt_ 
#endif 
#ifndef _In_ 
#define _In_ 
#endif 
#ifndef _Out_ 
#define _Out_ 
#endif 
#ifndef _Inout_ 
#define _Inout_
#endif

运行报错—The given version [14] is not supported, only version 1 to 10 is supported in this build.

将onnxruntime-1.14.0\onnxruntime-win-x64-1.14.0\lib的onnxruntime.dll复制一份到exe的目录下面，这是因为路径默认索引的是System32中的老版本库文件：

System32中存在老版本的onnx动态库：