如何在 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 上运行？我如何确认它正在工作？

Answer 1

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

然后：

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

Answer 2

get_device() 命令为您提供 onnxruntime 支持的设备。对于 CPU 和 GPU，有不同的运行时包可用。

目前您的onnxruntime环境仅支持CPU，因为您安装了CPU版本的onnxruntime。

如果您想构建供 GPU 使用的 onnxruntime 环境，请执行以下简单步骤。

第 1 步：卸载当前的 onnxruntime

>> pip uninstall onnxruntime

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

>>pip install onnxruntime-gpu

步骤 3：验证设备对 onnxruntime 环境的支持

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

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