同项目PyTorch可用CUDA，ONNX Runtime报CUDNN_STATUS_NOT_INITIALIZED错误

针对你遇到的PyTorch可正常调用CUDA，但ONNX Runtime使用CUDA推理时抛出CUDNN_STATUS_NOT_INITIALIZED错误的情况，可尝试以下几种解决方案：

1. 确保ONNX Runtime GPU版本与CUDA/cuDNN版本匹配 #

首先确认安装的onnxruntime-gpu版本与你的CUDA 11.8、cuDNN 8.7兼容。可重新安装对应版本的ONNX Runtime：

pip uninstall -y onnxruntime onnxruntime-gpu
pip install onnxruntime-gpu==1.15.1  # 该版本适配CUDA 11.8与cuDNN 8.7

2. 显式初始化CUDA上下文 #

PyTorch已初始化CUDA上下文，但ONNX Runtime可能无法正确复用，可在代码开头添加CUDA初始化逻辑：

import torch
# 显式初始化CUDA上下文
torch.cuda.init()
# 执行一个简单的CUDA运算确保上下文激活
_ = torch.tensor([1.0]).cuda()

3. 调整ONNX Runtime的CUDA Provider配置 #

在创建InferenceSession时，显式指定设备ID并配置cudnn参数：

providers = [
    ('CUDAExecutionProvider', {
        'device_id': 0,  # 指定使用的GPU设备ID（默认0）
        'cudnn_conv_algo_search': 'DEFAULT',  # 采用默认的cudnn卷积算法搜索策略
    })
]

4. 确保环境变量正确指向CUDA/cuDNN路径 #

若CUDA/cuDNN为自定义安装，需在代码开头设置环境变量：

import os
# 替换为你的CUDA安装路径，系统默认安装通常为/usr/local/cuda
cuda_path = os.path.expanduser("~/cuda")
os.environ['LD_LIBRARY_PATH'] = f"{cuda_path}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
os.environ['CUDNN_INCLUDE'] = f"{cuda_path}/include"

修改后的完整代码示例 #

import os.path
import torch
import onnxruntime
import numpy as np
import cv2
from onnxruntime import SessionOptions

# 显式初始化CUDA上下文
torch.cuda.init()
_ = torch.tensor([1.0]).cuda()

# 设置CUDA环境变量（根据实际安装路径调整）
cuda_path = os.path.expanduser("~/cuda")
os.environ['LD_LIBRARY_PATH'] = f"{cuda_path}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
os.environ['CUDNN_INCLUDE'] = f"{cuda_path}/include"

# Load the ONNX model
d = os.path.abspath(f"{__file__}/..")

onnx_model_path = f'{d}/RealESRGAN_x4plus_anime_6B.onnx'
# Specify CUDA execution provider with explicit config
providers = [
    ('CUDAExecutionProvider', {
        'device_id': 0,
        'cudnn_conv_algo_search': 'DEFAULT',
    })
]
sess_opt = SessionOptions()
ort_session = onnxruntime.InferenceSession(onnx_model_path, sess_opt, providers=providers)

# Load and preprocess input image
input_image_path = f"{d}/big.jpg"
input_image = cv2.imread(input_image_path)
input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
input_image = input_image.astype(np.float32) / 255.0  # Normalize to [0, 1]
input_image = np.transpose(input_image, (2, 0, 1))  # Change the shape from HWC to CHW
input_image = np.expand_dims(input_image, axis=0)  # Add batch dimension

# Perform inference
ort_inputs = {ort_session.get_inputs()[0].name: input_image}
ort_outs = ort_session.run(None, ort_inputs)

# Post-process the output
output_image = ort_outs[0][0]  # Assuming only one output
output_image = np.transpose(output_image, (1, 2, 0))  # Change the shape from CHW to HWC
output_image = (output_image * 255).clip(0, 255).astype(np.uint8)  # Convert back to uint8 and clip values

# Display or save the output image
out = f"{d}/out.jpg"  # Adjust output path as needed
cv2.imwrite(out, cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR))
print(f"out {out}")

内容的提问来源于stack exchange，提问作者chikadance