您如何在GPU上运行ONNX模型?

时间:2020-10-20 19:23:33

标签: deep-learning onnx onnxruntime onnx-coreml

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

2 个答案:

答案 0 :(得分:0)

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

然后:

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

答案 1 :(得分:0)

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

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

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

第一步:卸载你当前的onnxruntime

>> pip uninstall onnxruntime

第 2 步:安装 GPU 版本的 onnxruntime 环境

>>pip install onnxruntime-gpu

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

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

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

相关问题