onnx 1.16.0模型部署(python -- C++)

环境：

• anaconda

  • python 3.9

  • torch符合自己要求即可

  • onnx1.16.0

  • onnxruntime1.16.0（选择gpu版本，pip install onnxruntime-gpu==1.16.0）

• C++（vs2019）

  • onnxruntime1.16.0（下载版本与conda一致，也使用gpu版本）

  • opencv4.5.1（根据自己实际需求是否需要）VS2022+opevcv4.5.1-CSDN博客

---

第一步：将pth模型转为onnx

import torch
import torchvision.models as models
import onnx
from model import get_model
import os
import numpy as np
import onnxruntime as ort
​
def export_model_to_onnx(model, input_tensor, output_file):
    """
    将PyTorch模型转换为ONNX格式。
​
    参数:
    - model: PyTorch模型。
    - input_tensor: 用于模型推理的输入张量。
    - output_file: ONNX模型的输出文件路径。
    """
    # 设置模型为评估模式
    model.eval()
​
    # 导出模型
    torch.onnx.export(model,  # 运行模型
                      input_tensor,  # 模型输入 (或一个元组对多个输入)
                      output_file,  # 输出文件路径
                      export_params=True,  # 存储训练后的参数权重
                      opset_version=10,  # ONNX版本
                      do_constant_folding=True,  # 是否执行常量折叠优化
                      input_names=['input'],  # 输入名
                      output_names=['output'],  # 输出名
                      dynamic_axes={'input': {0: 'batch_size'},  # 输入的动态轴
                                    'output': {0: 'batch_size'}})  # 输出的动态轴
​
    print(f"Model has been converted to ONNX and saved as {output_file}.")
​
def test_onnx_model(onnx_path, input_tensor):
    """
    使用ONNX Runtime加载并运行ONNX模型，返回输出。
​
    参数:
    - onnx_path: ONNX模型文件的路径。
    - input_tensor: PyTorch张量，作为输入。
    """
    providers = ['CUDAExecutionProvider']  # 根据实际环境选择合适的执行提供者
    ort_session = ort.InferenceSession(onnx_path, providers=providers)
    ort_inputs = {ort_session.get_inputs()[0].name: input_tensor.numpy()}
    ort_outputs = ort_session.run(None, ort_inputs)
    return ort_outputs
​
​
# 主程序
if __name__ == '__main__':
    #----------------------------pth转onnx-----------------------------------#
    # 加载预训练模型，导入模型，根据你自己的模型加载
    model = get_model() 
​
    weights_path = "vgg3d_checkpoint.pth"
    assert os.path.exists(weights_path), f"file: '{weights_path}' does not exist."
    model.load_state_dict(torch.load(weights_path))
​
    # 创建一个示例输入（假设输入是一个1x42x42x42的3D图像）(batch_size, channel, w, h, d) ,根据实际调整，这里是三维数据，二维数据格式：(batch_size, channel, w, h)
    dummy_input = torch.randn(1, 1, 42, 42, 42, requires_grad=True)
​
    # 指定ONNX模型输出文件路径
    onnx_output_file = 'model.onnx'
​
    # 导出模型到ONNX
    export_model_to_onnx(model, dummy_input, onnx_output_file)
    # ----------------------------pth转onnx-----------------------------------#
​
    # ----------------------------验证onnx-----------------------------------#
    for i in range(2):
        # 创建输入数据
        dummy_input = torch.randn(10, 1, 42, 42, 42)
​
        # ONNX模型路径
        onnx_model_path = 'model.onnx'
​
        # 使用ONNX Runtime测试ONNX模型
        onnx_outputs = test_onnx_model(onnx_model_path, dummy_input)
​
        print("ONNX Model Output:", onnx_outputs)
    # ----------------------------验证onnx-----------------------------------#

---

第二步：将onnx模型部署C++

首先打开vs，创建一个空白项目，创建一个DeepModel.h头文件和DeepModel.cpp源文件

DeepModel.h文件添加以下代码:

#pragma once
#include <onnxruntime_cxx_api.h>
#include <iostream>
#include <assert.h>
#include <vector>
#include <numeric>
#include <cassert>
#include <memory>
#include <opencv2/core.hpp>
​
class DeepModel
{
public:
    DeepModel();
    ~DeepModel();
​
private:
    // 变量
    float* predict;     // 推理结果
​
    // 模型加载
    Ort::Session load_model(const std::wstring& model_path);
​
    // 模型推理
    void  model_inference(Ort::Session* session, const cv::Mat& input_image);
​
public: 
    // 获取模型
    Ort::Session GetModel(const std::wstring& model_path);
​
    // 获取推理结果
    float* GetPredict(Ort::Session* session, const cv::Mat& input_image);
};
​

DeepModel.cpp文件添加以下代码:

#include "DeepModel.h"
​
// 初始化
DeepModel::DeepModel()
{
}
​
// 析构
DeepModel::~DeepModel()
{
    predict = nullptr;
}
​
/****************************模型加载**************************************/
Ort::Session DeepModel::load_model(const std::wstring& model_path)
{
    try {
        Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNX_C_API");
        Ort::SessionOptions session_options;
       /* OrtCUDAProviderOptions cuda_options;
        cuda_options.device_id = 0;
        session_options.AppendExecutionProvider_CUDA(cuda_options);*/
        Ort::Session session(env, model_path.c_str(), session_options);
        return session;
    }
    catch (const std::exception& e) {
        std::cerr << "An error occurred: " << e.what() << std::endl;
        throw; // rethrow the exception to handle it in calling function
    }
}
// 模型推理
void DeepModel::model_inference(Ort::Session* session, const cv::Mat& input_image)
{
    /********** 预处理图片 **********/
    // 维度和数据按自己模型修改
    int batch_size = input_image.size[0];
    int int_samples_shape[5] = { batch_size, 1, 42, 42, 42 };
​
    size_t input_data_size = batch_size * 1 * 42 * 42 * 42;
    
    // 这个+5需要根据自己输入数据维度修改，batch_size, 1, 42, 42, 42 共五个维度
    std::vector<int64_t> samples_shape(int_samples_shape, int_samples_shape + 5);
​
    float* input_data = reinterpret_cast<float*>(input_image.data);
​
    /********** 从数据值创建输入张量对象 **********/
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info,
        input_data,
        input_data_size,
        samples_shape.data(),
        samples_shape.size()
        );
​
    /********** 设计输入和输出的名字 **********/
    Ort::AllocatorWithDefaultOptions allocator;
    std::vector<std::string> input_node_names;
    for (int i = 0; i < session->GetInputCount(); i++)
    {
        Ort::AllocatedStringPtr input_name = session->GetInputNameAllocated(i, allocator);
        input_node_names.push_back(input_name.get());  // 将其转换为 std::string 存储
    }
​
    std::vector<std::string> output_node_names;
    for (int i = 0; i < session->GetOutputCount(); i++)
    {
        Ort::AllocatedStringPtr output_name = session->GetOutputNameAllocated(i, allocator);
        output_node_names.push_back(output_name.get());  // 将其转换为 std::string 存储
    }
​
    std::vector<const char*> input_node_names_cstr;
    for (const auto& name : input_node_names)
    {
        input_node_names_cstr.push_back(name.c_str());  // 将 std::string 转换为 const char*
    }
​
    std::vector<const char*> output_node_names_cstr;
    for (const auto& name : output_node_names)
    {
        output_node_names_cstr.push_back(name.c_str());  // 将 std::string 转换为 const char*
    }
​
    /********** 输入图片数据，运行模型获取预测结果 **********/
    std::vector<Ort::Value> output_tensors = session->Run(
        Ort::RunOptions{ nullptr },
        input_node_names_cstr.data(),
        &input_tensor,
        1,
        output_node_names_cstr.data(),
        output_node_names_cstr.size()
    );
​
    predict = output_tensors[0].GetTensorMutableData<float>();
}
​
/****************************获取模型**************************************/
Ort::Session DeepModel::GetModel(const std::wstring& model_path)
{
    return load_model(model_path);
}
​
// 获取模型预测结果
float* DeepModel::GetPredict(Ort::Session* session, const cv::Mat& input_image)
{
    model_inference(session, input_image);
    return predict;
}

然后创建一个main.cpp文件，添加一下测试代码：

#pragma once
#include "DeepModel.h"
#include <Windows.h>
​
int main() {
    
    DeepModel model;
    try {
        // 加载模型，onnx模型路径放在项目文件夹下
        auto session = model.GetModel(L"model.onnx");
​
        if (slice_session) {
            std::cout << "slice model load success!" << std::endl;
        }
​
        if (chong_session) {
            std::cout << "chong model load success!" << std::endl;
        }
​
       int batch_size = 1;
        // 根据自己的输入修改
        cv::Mat input_image(batch_size, 1 * 42 * 42 * 42, CV_32F);
​
        input_image = cv::Scalar(1);
​
        float* predict = model.GetPredict(&session, input_image);
        
        // 根据自己的分类数量修改 batch_size * 2，我的是两个类别
        for (int j = 0; j < batch_size * 2; ++j) {
            std::cout << j << "::" << predict[j] << std::endl;
        }
        predict = nullptr;
        return 0;
    }
    catch (const std::exception& e) {
        std::cerr << "An error occurred: " << e.what() << std::endl;
    }
​
    return 0;
}

第三步：配置vs的配置（release x64）根据自己需求修改，Debug配置一样

• VC++目录

到这里以及完成。运行代码即可。