本文最后更新于:2025年1月2日 下午

在 onnxruntime python 库中可以加载在内存或者显存中,分别使用 cpu 和 gpu 进行推断运算,本文记录 onnxruntime 将模型加载到指定 gpu 的步骤。

安装

这里我们使用 1.18.0 的 onnx 库

1
pip install onnxruntime-gpu==1.18.0

指定 gpu 配置

将需要加载模型的 gpu id 填入 device_id,需要用字符串类型

1
2
3
4
5
6
7
8
9
import onnxruntime
model_path = 'test.onnx'
CUDAExecutionProvider = {
"cudnn_conv_algo_search": "DEFAULT",
"cudnn_conv_use_max_workspace": '1',
"device_id": '1'
}
providers=[("CUDAExecutionProvider", CUDAExecutionProvider)]
session = onnxruntime.InferenceSession(model_path, providers)

也有一种实现思路是在 Docker 内仅映射部分 gpu

1
docker --gpus 'devices=1' 

参考资料



文章链接:
https://www.zywvvd.com/notes/study/deep-learning/deploy/onnx-percific-gpu/onnx-percific-gpu/


“觉得不错的话,给点打赏吧 ୧(๑•̀⌄•́๑)૭”

微信二维码

微信支付

支付宝二维码

支付宝支付

Python onnxruntime-gpu 指定 gpu 加载模型
https://www.zywvvd.com/notes/study/deep-learning/deploy/onnx-percific-gpu/onnx-percific-gpu/
作者
Yiwei Zhang
发布于
2025年1月2日
许可协议