如何使用 ONNX 结合 GPU 加速推理(CUDA 与 cuDNN 简明指南)

发布于:2024-09-18 ⋅ 阅读:(114) ⋅ 点赞:(0)

前言

在深度学习模型推理中,使用 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 进行硬件加速都能显著提升模型的执行效率。

分享完成~