pytorch转onnx

将 PyTorch 模型转换为 ONNX 格式，可以使用以下步骤：

1. 首先，安装好 PyTorch 和 ONNX 的 Python 包。

2. 加载 PyTorch 模型，并将其转换为 ONNX 格式：

import torch
import onnx

# 加载 PyTorch 模型
model = torch.load('model.pt')

# 转换为 ONNX 格式
dummy_input = torch.randn(1, 3, 224, 224) # 示例输入
input_names = ['input']
output_names = ['output']
onnx_model = onnx.export(model, dummy_input, 'model.onnx', input_names=input_names, output_names=output_names)

3. 将 ONNX 模型加载到 C++ 程序中：

#include <iostream>
#include <fstream>
#include <string>
#include <vector>

#include <onnxruntime/core/session/onnxruntime_cxx_api.h>

int main() {
    // 加载 ONNX 模型
    const char* model_path = "model.onnx";
    Ort::SessionOptions session_options;
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    Ort::Session session(env, model_path, session_options);

    // 构建输入张量
    std::vector<float> input_tensor_values(3 * 224 * 224, 1.0f); // 示例输入数据
    std::vector<int64_t> input_tensor_shape{1, 3, 224, 224};
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_tensor_values.data(), input_tensor_values.size(), input_tensor_shape.data(), input_tensor_shape.size());

    // 运行模型
    std::vector<const char*> input_names{"input"};
    std::vector<const char*> output_names{"output"};
    std::vector<Ort::Value> output_tensors = session.Run(Ort::RunOptions{nullptr}, input_names.data(), &input_tensor, 1, output_names.data(), 1);
    Ort::Value& output_tensor = output_tensors.front();

    // 获取输出张量
    float* output_tensor_data = output_tensor.GetTensorMutableData<float>();
    std::vector<int64_t> output_tensor_shape(output_tensor.GetTensorTypeAndShapeInfo().GetShape());
    std::cout << "Output tensor shape: ";
    for (const auto& dim : output_tensor_shape) {
        std::cout << dim << " ";
    }
    std::cout << std::endl;

    return 0;
}

这样就可以将 PyTorch 模型转换为 ONNX 格式，并在 C++ 程序中加载和运行模型了。