CUDA Execution Provider

CUDA 执行引擎使得在支持 Nvidia CUDA 技术的 GPU 上能够实现硬件加速计算。

onnxruntime-gpu 包的设计理念是能够与 PyTorch 无缝配合使用，前提是两者都是基于相同版本的 CUDA 和 cuDNN 构建的。在安装支持 CUDA 功能的 PyTorch 时，系统中会自动包含所需的 CUDA 和 cuDNN DLL 文件，因此无需单独安装 CUDA 工具包或 cuDNN。

为确保 ONNX Runtime 能够使用 PyTorch 所安装的 DLL 文件，你可以在创建推理会话之前预先加载这些库。实现方式可以是直接导入 PyTorch，或者使用 onnxruntime.preload_dlls() 函数。

import onnxruntime as ort

model_path = '<path to model>'

providers = [
    ('CUDAExecutionProvider', {
        'device_id': 0,
        'arena_extend_strategy': 'kNextPowerOfTwo',
        'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
        'cudnn_conv_algo_search': 'EXHAUSTIVE',
        'do_copy_in_default_stream': True,
    }),
    'CPUExecutionProvider',
]

session = ort.InferenceSession(model_path, providers=providers)

Configuration Options   配置选项