onnxruntime-gpu使用记录

1 onnxruntime-gpu的优势

相比于onnxruntime,使用onnxruntime-gpu在gpu上进行模型推理的速度会明显加快。

2 onnxruntime-gpu的安装

如果已经安装了onnxruntime或者同时安装了onnxruntime和onnxruntime-gpu,需先卸载掉这些包,

pip uninstall onnxruntime
or
pip uninstall onnxruntime
pip uninstall onnxruntime-gpu

接下来,安装onnxruntime-gpu,在安装前,需要查看cuda版本,再根据cuda的版本安装相应的onnxruntime-gpu包,

# 查看cuda版本
nvidia-smi
onnxruntime-gpu使用记录_第1张图片

去官网查看cuda跟onnxruntime的对应版本,https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html

onnxruntime-gpu使用记录_第2张图片

安装onnxruntime-gpu,

pip install onnxruntime-gpu==1.9

3 检验onnxruntime-gpu是否能正常使用

import onnxruntime as ort

print(ort.get_device())
ort_session = ort.InferenceSession('/kaxier01/projects/230203_sim.onnx', providers=['CUDAExecutionProvider'])
print(ort_session.get_providers())


## 结果
GPU
['CUDAExecutionProvider', 'CPUExecutionProvider']

你可能感兴趣的:(python,计算机视觉,人工智能,onnx)