要在 C++ 中加载并推理 PyTorch 训练好的模型，通常会借助 ONNX 作为中间格式。以下是具体方法： ### 使用 ONNX 转换和加载 #### 将 PyTorch 模型导出为 ONNX 格式 首先，在 Python 环境中使用 `torch.onnx.export` 方法将 PyTorch 模型保存为 ONNX 文件。 ```python import torch import torchvision.models as models # 加载预训练模型 model = models.resnet18(pretrained=True) model.eval() # 创建虚拟输入张量 dummy_input = torch.randn(1, 3, 224, 224) # 导出为 ONNX 格式 torch.onnx.export(model, dummy_input, "resnet18.onnx", input_names=['input'], output_names=['output']) ``` 此过程将生成名为 `resnet18.onnx` 的文件[^3]。 --- #### 在 C++ 中加载 ONNX 模型 接下来，利用 ONNX Runtime 或其他支持 ONNX 的工具包（如 TensorRT）在 C++ 中加载该模型。 ##### 安装依赖项 确保安装了 ONNX Runtime 库及其头文件。可以通过以下命令安装： ```bash sudo apt-get install libonnxruntime-dev ``` 或者下载源码编译安装。 --- ##### 编写 C++ 推理代码 下面展示了一个完整的例子，说明如何通过 ONNX Runtime 进行推断。 ```cpp #include <onnxruntime_cxx_api.h> #include <iostream> #include <vector> int main() { Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test"); // 初始化环境 Ort::SessionOptions session_options; Ort::Session session(env, "resnet18.onnx", session_options); // 获取输入节点名称 auto input_node_names = std::vector<const char *>{session.GetInputNameAllocated(0, nullptr).get()}; // 设置输入数据 (假设输入形状为 [1, 3, 224, 224]) const float data[] = { /* 填充实际数据 */ }; size_t tensor_size = sizeof(data) / sizeof(float); Ort::Value input_tensor = Ort::Value::CreateTensor<float>( Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault), data, tensor_size, {1, 3, 224, 224}, 4 ); // 输出节点名 auto output_node_names = std::vector<const char *>{session.GetOutputNameAllocated(0, nullptr).get()}; // 执行推理 auto output_tensors = session.Run( Ort::RunOptions{nullptr}, input_node_names.data(), &input_tensor, 1, output_node_names.data(), 1 ); // 处理输出 float* output_data = output_tensors[0].GetTensorMutableData<float>(); for (size_t i = 0; i < 10; ++i) { std::cout << output_data[i] << ", "; } } ``` 上述代码展示了如何初始化 ONNX Runtime 并运行推理操作。 --- ### 替代方案：使用 TensorRT 如果性能至关重要，则可以选择 TensorRT 来优化 ONNX 模型的执行效率。其流程如下： 1. **导入 ONNX 模型到 TensorRT** 利用 TensorRT 提供的解析器读取 ONNX 文件。 2. **构建引擎** 构建序列化后的 TensorRT 引擎用于高效推理。 示例代码片段： ```cpp #include <NvInfer.h> #include <onnx_parser.h> #include <fstream> nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger.getTRTLogger()); auto network = builder->createNetworkV2(0U); auto parser = nvonnxparser::createParser(*network, gLogger); // 解析 ONNX 文件 if (!parser->parseFromFile("resnet18.onnx", static_cast<int>(gLogger Severity.Info))) { throw std::runtime_error("Failed to parse the ONNX file"); } // 配置并构建引擎 builder->setMaxBatchSize(1); auto config = builder->createBuilderConfig(); config->setFlag(nvinfer1::BuilderFlag::kFP16); // 启用 FP16 如果硬件支持 auto engine = builder->buildEngineWithConfig(*network, *config); ``` 以上代码实现了从 ONNX 到 TensorRT 引擎的转换[^1]。 --- ### 总结 无论是采用 ONNX Runtime 还是 TensorRT，都需要先将 PyTorch 模型导出为 ONNX 格式。之后可根据需求选择适合的框架完成后续工作。