GPU 依赖项
对于 GPU 版本的 ONNX Runtime,需要安装 CUDA 和 cuDNN。请检查 CUDA 执行提供程序的要求以获取兼容版本的 CUDA 和 cuDNN。安装时请注意:
- cuDNN 需要 ZLib,安装方法请参照 cuDNN 安装指南。
- CUDA 和 cuDNN 的 bin 目录路径需要添加到环境变量
PATH
中。
《这部分可以去B站有很多教学视频讲的更详细》
Python 环境安装
安装 ONNX Runtime CPU 版本
要在 Python 环境中安装 ONNX Runtime 的 CPU 版本,可以使用以下命令:
pip install onnxruntime
安装 ONNX Runtime GPU 版本 (CUDA )
pip install onnxruntime-gpu
验证Python 示例
以下是一个简单的 Python 示例,用于加载和运行 ONNX 模型:
import onnxruntime as ort
# 加载 ONNX 模型
session = ort.InferenceSession("model.onnx")
# 准备输入数据
input_name = session.get_inputs()[0].name
input_data = ... # 根据模型要求准备输入数据
# 运行推理
result = session.run(None, {input_name: input_data})
# 输出结果
print(result)
C++ 环境安装
安装 ONNX Runtime CPU 版本
在 C++ 环境中安装 ONNX Runtime 的 CPU 版本,可以使用以下命令:
# 克隆 ONNX Runtime 仓库
git clone --recursive https://github.com/microsoft/onnxruntime
cd onnxruntime
# 构建并安装
./build.sh --config Release --build_shared_lib
安装 ONNX Runtime GPU 版本 (CUDA 11.x)
要安装支持 CUDA 11.x 的 GPU 版本,可以使用以下步骤:
-
设置环境变量
export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
2. 构建 ONNX Runtime
./build.sh --config Release --use_cuda
《windows端可以用 vs编辑器配置相应的头文件、库文件》
C++ 验证示例
以下是一个简单的 C++ 示例,用于加载和运行 ONNX 模型:
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
#include <vector>
#include <iostream>
int main() {
// 初始化 ONNX Runtime
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "example");
// 创建 SessionOptions
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);
// 加载模型
Ort::Session session(env, "model.onnx", session_options);
// 获取模型输入信息
Ort::AllocatorWithDefaultOptions allocator;
const char* input_name = session.GetInputName(0, allocator);
std::cout << "Input Name: " << input_name << std::endl;
// 准备输入数据
std::vector<float> input_data = ...; // 根据模型要求准备输入数据
std::vector<int64_t> input_shape = {1, 3, 224, 224}; // 示例输入形状
// 创建输入 tensor
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(allocator, input_data.data(), input_data.size(), input_shape.data(), input_shape.size());
// 运行推理
auto output_tensors = session.Run(Ort::RunOptions{nullptr}, &input_name, &input_tensor, 1, session.GetOutputNames(allocator), 1);
// 输出结果
float* output_data = output_tensors[0].GetTensorMutableData<float>();
std::cout << "Output: " << output_data[0] << std::endl;
return 0;
}
通过上述步骤和示例代码,您可以在 Python 和 C++ 环境下安装并使用 ONNX Runtime 进行模型推理。