如何在 GPU 上运行 ONNX 模型？答案

【问题标题】：How do you run a ONNX model on a GPU?如何在 GPU 上运行 ONNX 模型？
【发布时间】：2021-02-03 16:27:54
【问题描述】：

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

【问题讨论】：

标签： deep-learning onnx onnxruntime onnx-coreml

【解决方案1】：

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

然后：

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

【讨论】：

【解决方案2】：

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

目前你的onnxruntime环境只支持CPU，因为你已经安装了onnxruntime的CPU版本。

如果您想为 GPU 构建 onnxruntime 环境，请按照简单的步骤操作。

第 1 步：卸载您当前的 onnxruntime

>> pip uninstall onnxruntime

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

>>pip install onnxruntime-gpu

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

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

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

【讨论】：

【解决方案3】：

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

【讨论】：