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

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

安装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']
GitHub 加速计划 / on / onnxruntime
13.76 K
2.79 K
下载
microsoft/onnxruntime: 是一个用于运行各种机器学习模型的开源库。适合对机器学习和深度学习有兴趣的人,特别是在开发和部署机器学习模型时需要处理各种不同框架和算子的人。特点是支持多种机器学习框架和算子,包括 TensorFlow、PyTorch、Caffe 等,具有高性能和广泛的兼容性。
最近提交(Master分支:1 个月前 )
33217350 follow up to https://github.com/microsoft/onnxruntime/pull/22282 replaces https://github.com/microsoft/onnxruntime/pull/22388 15 小时前
18e81f87 ### Description This PR fixes a bug when the ONNX checker is called while exporting Whisper for FP16 CUDA with optional flags. ### Motivation and Context Sometimes, the ONNX checker raises an error depending on the optional flags passed. By wrapping the ONNX checker in a try-except, the conversion can continue even if the checker fails. 19 小时前
Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐