版本
onnxruntime-gpu版本可以说是一个非常简单易用的框架,因为通常用pytorch训练的模型,在部署时,会首先转换成onnx,而onnxruntime和onnx又是有着同一个爸爸,无疑,在op的支持上肯定是最好的。
通常在安装onnxruntime时,需要将其版本与pytorch版本和CUDA版本进行对应,其中ONNXRuntime与CUDA版本对应关系表如下表所示。
选好版本后,使用以下命令进行安装onnxruntime
pip install onnxruntime-gpu==X.X.X -i https://pypi.tuna.tsinghua.edu.cn/simple
查看是否使用上GPU
python 3.8.17 (default, Jul 5 2023, 21:04:15)
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import onnxruntime
>>> onnxruntime.get_device()
'GPU'
>>> onnxruntime.get_available_providers()
['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'AzureExecutionProvider', 'CPUExecutionProvider']