果然一知半解是不得行的!
- 安装onnxruntime-gpu(注意是gpu版本)
pip install onnxruntime-gpu
- 确认是否安装成功
>>> import onnxruntime
>>> onnxruntime.get_device()
'GPU'
>>> onnxruntime.get_available_providers()
['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
>>> exit()
- 部署
新建InferenceSession的时候,加入 TensorrtExecutionProvider 和 CUDAExecutionProvider 就可以了。
值得注意的是,‘provider’,很多源码在使用CPU是设置为None了,要改成一下形式CPU和GPU部署都能通用。
session = onnxruntime.InferenceSession(
"YOUR-ONNX-MODEL-PATH",
providers=onnxruntime.get_available_providers()
)