安装onnxruntime 推理加速

Posted 2022-05-25 qianbo_insist

1、临时更改源

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnxruntime

如果是默认的源位置会很慢，更改以下为好

2、测试

cmd下键入python

 python>>> import onnxruntime
 python>>> onnxruntime.get_device()
 'CPU'
 python>>>  onnxruntime.get_available_providers()
['CPUExecutionProvider']

这样安装仅仅是cpu版本,只支持cpu模式，下面删除后安装gpu版本

pip uninstall onnxruntime

删除后重新安装

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnxruntime-gpu

Python 3.7.2 (tags/v3.7.2:9a3ffc0492, Dec 23 2018, 23:09:28) [MSC v.1916 64 bit (AMD64)] on win32
Type “help”, “copyright”, “credits” or “license” for more information.

 import onnxruntime
 onnxruntime.get_device()
'GPU'
 onnxruntime.get_available_providers()
['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']

可见除了cpu方式，另外支持了TensorRt 和 CUDA 模式

例子

import onnxruntime
import numpy as np

#device_name = 'GPU' 
device_name = onnxruntime.get_device()

if device_name == 'cpu':
    providers = ['CPUExecutionProvider']
elif device_name == 'GPU':
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
onnx_model = onnxruntime.InferenceSession('模型名称.onnx', providers=providers)
#这里输入数据进行推理
data = .......
onnx_input = onnx_model.get_inputs()[0].name: data
outputs = onnx_model.run(None, onnx_input)