当前位置：首页 > news >正文

openvino如何在c++中调用pytorch训练的模型

news 来源：原创 2025/6/7 10:58:48

步骤1：将PyTorch模型转换为ONNX格式

转换代码示例（Python）

import torch
import torchvision

1. 加载训练好的PyTorch模型

model = torchvision.models.resnet18(pretrained=True)
model.eval()  # 设置为评估模式

2. 创建虚拟输入（尺寸需匹配模型输入）

dummy_input = torch.randn(1, 3, 224, 224)

3. 导出为ONNX格式

torch.onnx.export(model,dummy_input,"resnet18.onnx",export_params=True,opset_version=11,  # 推荐使用11+input_names=["input"],output_names=["output"]
)

步骤2：将ONNX转换为OpenVINO IR格式

使用OpenVINO的Model Optimizer转换模型：

mo --input_model resnet18.onnx \--output_dir ir_model \--model_name resnet18_ir \--data_type FP32  # 也可用FP16

生成文件：
resnet18_ir.xml：网络拓扑结构
resnet18_ir.bin：权重数据

步骤3：C++集成推理代码

环境配置
安装OpenVINO Runtime
CMake配置示例：

find_package(OpenVINO REQUIRED)
add_executable(inference_app main.cpp)
target_link_libraries(inference_app PRIVATE openvino::runtime)

C++推理代码

#include <openvino/openvino.hpp>
#include <vector>int main() {// 1. 初始化OpenVINO核心ov::Core core;// 2. 加载模型auto model = core.read_model("ir_model/resnet18_ir.xml");// 3. 编译模型（指定设备）ov::CompiledModel compiled_model = core.compile_model(model, "CPU");  // 也可用"GPU", "MYRIAD"// 4. 创建推理请求ov::InferRequest infer_request = compiled_model.create_infer_request();// 5. 准备输入数据auto input_tensor = infer_request.get_input_tensor();float* input_data = input_tensor.data<float>();// 填充数据（示例：随机值）std::fill_n(input_data, input_tensor.get_size(), 0.5f); // 6. 执行推理infer_request.infer();// 7. 获取输出结果auto output_tensor = infer_request.get_output_tensor();const float* output_data = output_tensor.data<const float>();// 处理输出（例如：打印前10个结果）for (int i = 0; i < 10; ++i) {std::cout << "Output[" << i << "] = " << output_data[i] << std::endl;}return 0;
}

关键注意事项

输入预处理：
需在C++中复现PyTorch的预处理逻辑（归一化/缩放）
示例：若PyTorch使用mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]，需在填充input_data前进行归一化
动态shape支持：
若需动态输入尺寸，导出ONNX时指定动态维度：
torch.onnx.export(…, dynamic_axes={“input”: {0: “batch_size”, 2: “height”, 3: “width”}})

性能优化：
使用ov::preprocess::PrePostProcessor进行图内预处理
启用异步推理：infer_request.start_async()
多设备支持：core.compile_model(model, “MULTI:CPU,GPU”)
错误排查工具：
验证ONNX：onnx.checker.check_model(onnx.load(“resnet18.onnx”))
查看IR结构：使用Netron打开.xml文件

完整流程图示

graph LR
A[PyTorch模型 .pt] -->|torch.onnx.export| B[ONNX模型]
B -->|Model Optimizer| C[OpenVINO IR XML/BIN]
C --> D[C++加载IR]
D --> E[预处理数据]
E --> F[执行推理]
F --> G[解析输出]

通过以上步骤，即可高效部署PyTorch模型到C++生产环境。建议参考OpenVINO官方文档获取最新API细节和性能调优指南。