ONNX Runtime GPU无法调用求助：WD14模型仅运行于CPU

问题背景 #

使用imgutils库的get_wd14_tags函数基于WD14模型做图像标注，代码功能正常但仅使用CPU；系统中PyTorch、Stable Diffusion等其他AI应用可正常调用GPU。已确认：

• 安装了onnxruntime-gpu==1.16.2，ort.get_all_providers()显示CUDAExecutionProvider可用
• ONNX Runtime初始化会话时providers列表包含['CUDAExecutionProvider', 'CPUExecutionProvider']，但强制修改provider无效
• 环境：Debian Bookworm、NVIDIA驱动545.23.06、CUDA 12.3

具体解决思路 #

1. 验证ONNX Runtime与CUDA版本兼容性 #

ONNX Runtime对CUDA版本有严格适配要求：onnxruntime-gpu==1.16.2官方推荐搭配CUDA 11.8/12.1，当前使用的CUDA 12.3不在兼容范围内，可能导致CUDA provider无法正常加载。

• 方案：降级CUDA到12.1，或升级onnxruntime-gpu到1.17及以上版本（该版本开始支持CUDA 12.3）

2. 绕过封装函数，手动指定Provider初始化模型 #

imgutils的get_wd14_tags可能封装了模型初始化逻辑，未暴露provider配置项。可手动加载WD14的ONNX模型并强制指定CUDA provider：

import glob
import numpy as np
import onnxruntime as ort
from imgutils.data import load_image_pil
from imgutils.tagging.wd14 import _parse_wd14_tags  # 复用imgutils的标签解析逻辑

# 替换为你实际的WD14模型路径（可通过imgutils源码找到默认模型路径）
model_path = "/path/to/wd14_model.onnx"

# 初始化会话时强制指定CUDA provider
sess_options = ort.SessionOptions()
session = ort.InferenceSession(model_path, sess_options, providers=['CUDAExecutionProvider'])

# 遍历图片处理
for image in glob.glob('/picture_store/*.JPG'):
    # 按WD14要求预处理图像
    img = load_image_pil(image, mode='RGB', size=(448, 448))
    img_tensor = np.array(img).transpose(2, 0, 1).astype(np.float32) / 255.0
    img_tensor = np.expand_dims(img_tensor, axis=0)
    
    # 模型推理
    outputs = session.run(None, {'input': img_tensor})
    rating, features, chars = _parse_wd14_tags(outputs[0], outputs[1])
    
    print(image, rating, features, chars)

3. 验证ONNX Runtime GPU基础功能 #

编写最小测试脚本，确认onnxruntime-gpu能实际调用GPU：

import onnxruntime as ort
import numpy as np

# 用WD14模型或导出的简单测试模型测试
sess = ort.InferenceSession('dummy_model.onnx', providers=['CUDAExecutionProvider'])
print(f"配置的Providers: {sess.get_providers()}")
print(f"实际使用的Providers: {sess.get_providers()}")

# 运行推理
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = sess.run(None, {'input': input_data})

如果测试脚本也无法使用GPU，说明onnxruntime-gpu本身的GPU支持存在问题，需重新安装匹配版本的依赖。

4. 检查环境变量与权限 #

• 确认当前用户可通过nvidia-smi正常查看GPU状态，无权限问题
• 检查环境变量CUDA_PATH、LD_LIBRARY_PATH是否正确指向CUDA 12.3的安装路径，避免onnxruntime加载到系统默认的CPU版本库
• 确保没有冲突的onnxruntime版本（如同时安装了onnxruntime和onnxruntime-gpu），可执行pip uninstall onnxruntime移除CPU版本

5. 开启ONNX Runtime调试日志定位问题 #

开启DEBUG级日志，查看会话初始化时CUDA加载的具体错误信息：

import onnxruntime as ort

# 设置日志级别为DEBUG（0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR）
ort.set_default_logger_severity(0)

# 运行原图像标注代码
import glob
from imgutils.tagging import get_wd14_tags

for image in glob.glob('/picture_store/*.JPG'):
    rating, features, chars = get_wd14_tags(image)
    print(image, rating, features, chars)

日志中会包含CUDA provider初始化失败的具体原因（如缺失cudnn库、版本不匹配等）。

内容的提问来源于stack exchange，提问作者Dr. Volker Jaenisch