前言
在深度学习模型推理中,使用 GPU 进行加速是提升模型推理速度的关键方式之一。
本文将带大家一步步了解如何使用 ONNX Runtime 结合 NVIDIA 的 CUDA 和 cuDNN 进行 GPU 加速。
一、查找ONNX、CUDA与cuDNN之间的对应版本
首先,我们需要确保 ONNX Runtime 与 CUDA 和 cuDNN 的版本兼容。
如果版本不匹配,可能导致加速失败或性能下降。我们可以通过官方兼容性文档查找它们的对应关系。
来到ONNX官网地址查询:https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html
这里有详细版本对应关系,比如 CUDA12.x版本的:
比如 CUDA11.x版本的,这里只展示了部分;后还有CUDA10.x等版本的对应关系。
二、安装onnxruntime-gpu
接下来,我们需要安装支持 GPU 的 ONNX Runtime。默认安装的 onnxruntime
只支持CPU
如果我们想要使用 GPU(CUDA),需要安装支持 CUDA 的onnxruntime-gpu版本。
pip install onnxruntime-gpu
需要根据上面查询到的,指定onnxruntime-gpu版本,比如1.19.0版本,并指定使用清华源加速
pip install onnxruntime-gpu==1.19.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
温馨提示:onnxruntime-gpu 1.19.0版本,对应CUDA12.2,cuDNN9.3
三、安装CUDA与cuDNN
安装完 onnxruntime-gpu
后,下一步是确保 CUDA 和 cuDNN 库已正确安装,并且版本与 ONNX Runtime 兼容。
这里建议参考博客:https://blog.csdn.net/YYDS_WV/article/details/137825313
CUDA的官网地址:https://developer.nvidia.com/cuda-toolkit-archive
cuDNN的官网地址:https://developer.nvidia.com/rdp/cudnn-archive
下载并安装这些工具后,确保将其路径正确添加到系统的环境变量中,方便 ONNX Runtime 使用。
安装完成后,打开Windows系统的命令终端
用nvcc -V命令查询安装信息,看到对应CUDA版本
若无法返回结果,则表明CUDA安装失败,或环境变量没有设置对。
PS C:\Windows\system32> nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Jun_13_19:42:34_Pacific_Daylight_Time_2023
Cuda compilation tools, release 12.2, V12.2.91
Build cuda_12.2.r12.2/compiler.32965470_0
PS C:\Windows\system32>
四、检查ONNX加速是否成功
在所有软件环境配置完成后,最后一步是检查 ONNX 是否成功启用了 GPU 加速。
1、写一个简单的测试代码,使用onnxruntime
加载ONNX模型,并通过CUDA进行推理,同时检查是否成功使用了CUDA加速。
import onnxruntime as ort
# 创建一个推理会话
session = ort.InferenceSession(r"YOLO_t22_best.onnx", providers=['CUDAExecutionProvider'])
# 检查是否使用了CUDA
providers = session.get_providers()
print(f"Available providers: {providers}")
# 获取当前执行程序的是否使用GPU设备
device = ort.get_device()
print(f"Current device: {device}")
查看返回信息,如果显示有CUDAExecutionProvider 和 GPU,说明使用了GPU加速推理
Available providers: ['CUDAExecutionProvider', 'CPUExecutionProvider']
Current device: GPU
如果没有安装 CUDA,或者CUDA未正确配置,以上代码的输出只有 CPUExecutionProvider
2、或者直接检查 CUDA 是否可用
import onnxruntime as ort
# 检查是否有CUDA加速
available_providers = ort.get_available_providers()
print(f"Available providers: {available_providers}")
# 检查当前设备是GPU还是CPU
device = ort.get_device()
print(f"Current device: {device}")
可以创建一个空的推理会话,或者仅使用onnxruntime
提供的 API来检查是否可以使用 CUDA,而无需加载任何模型。
查看返回信息,如果显示有CUDAExecutionProvider 和 GPU,说明使用了GPU加速推理
vailable providers: ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
Current device: GPU
3、编写一个程序,加载ONNX模型,对比GPU加速与CPU版本的速度
我们可以自定义运行次数,默认值为1000次
import onnxruntime as ort
import numpy as np
import time
# 生成随机输入数据(根据模型的输入要求调整大小)
input_shape = (1, 3, 640, 640) # 假设模型需要 640x640 大小的图像输入
input_data = np.random.rand(*input_shape).astype(np.float32)
# 加载 ONNX 模型路径
model_path = r"YOLO_t22_best.onnx"
# 自定义运行次数,默认为 1000
def benchmark(session, input_data, iterations=1000):
start_time = time.time()
for _ in range(iterations):
session.run(None, {session.get_inputs()[0].name: input_data})
return time.time() - start_time
# 1. 使用 CPU 进行推理
session_cpu = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
cpu_iterations = 1000 # 可修改为自定义次数
cpu_time = benchmark(session_cpu, input_data, iterations=cpu_iterations)
print(f"CPU 推理总时间: {cpu_time:.4f} 秒, 每次推理平均时间: {cpu_time / cpu_iterations:.4f} 秒")
# 2. 使用 GPU (CUDA) 进行推理
session_gpu = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider'])
gpu_iterations = 1000 # 可修改为自定义次数
gpu_time = benchmark(session_gpu, input_data, iterations=gpu_iterations)
print(f"GPU 推理总时间: {gpu_time:.4f} 秒, 每次推理平均时间: {gpu_time / gpu_iterations:.4f} 秒")
# 比较 GPU 和 CPU 的速度
speedup = cpu_time / gpu_time
print(f"GPU 加速比: {speedup:.2f} 倍")
计算并打印 CPU 和GPU 各自的总推理时间和每次推理的平均时间
我这个是老显卡了1080TI,加速也不是太明显;
如果想更准确,需要把迭代次数调大
CPU 推理总时间: 71.8592 秒, 每次推理平均时间: 0.1437 秒
GPU 推理总时间: 13.5614 秒, 每次推理平均时间: 0.0271 秒
GPU 加速比: 5.30 倍
通过以上步骤,我们可以轻松地在自己的系统中设置 ONNX Runtime 的 GPU 加速。
无论是用于训练模型还是在生产环境中进行推理,使用 CUDA 和 cuDNN 进行硬件加速都能显著提升模型的执行效率。
分享完成~