如何在GPU上运行ONNX模型？

Question

我正在尝试运行一个ONNX模型。

import onnxruntime as ort
import onnxruntime.backend
model_path = "model.onnx"

#https://microsoft.github.io/onnxruntime/
ort_sess = ort.InferenceSession(model_path)


print( ort.get_device()  )

这将打印出来

cpu

如何让它在我的GPU上运行？如何确认它正在工作？

- djacobs7

Answer 1

您可能安装了CPU版本。尝试卸载onnxruntime并安装GPU版本，例如pip install onnxruntime-gpu。

然后：

>>> import onnxruntime as ort
>>> ort.get_device()
'GPU'

- Sergii Dymchenko

---

单行版本将是 python -c "import onnxruntime as ort;print(ort.get_device())" - undefined

Answer 2

get_device()命令可以提供onnxruntime支持的设备信息。针对CPU和GPU，可用的运行时包是不同的。

目前，由于您安装的是CPU版本的onnxruntime，因此您的onnxruntime环境仅支持CPU。

如果您想要构建适用于GPU的onnxruntime环境，请按照以下简单步骤进行操作：

步骤1：卸载您当前的onnxruntime

>> pip uninstall onnxruntime

步骤二：安装 onnxruntime 环境的 GPU 版本

>>pip install onnxruntime-gpu

步骤3：验证设备是否支持onnxruntime环境

>> import onnxruntime as rt
>> rt.get_device()
'GPU'

步骤4：如果您遇到任何问题，请检查您的cuda和CuDNN版本，它们必须相互兼容。 请参考此链接here以了解cuda和CuDNN之间的版本兼容性。

- Abhijit Manepatil

Answer 3

你的onnxruntime-gpu版本应该与你的cuda和cudnn版本匹配，你可以从官方网站上查看它们之间的关系： https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html。

- bigbro

---

我认为这是我的问题所在。看起来目前还没有支持CUDA 11.6的版本。谢谢。 - Casey Jones