EADST

ONNX Set Device ID with CUDA

ONNX Set the Device ID with CUDA

import onnxruntime as ort

model_path = ''

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)

Reference:

CUDA Execution Provider

相关标签
About Me
XD
Goals determine what you are going to be.
Category
标签云
站点统计

本站现有博文240篇,共被浏览280696

本站已经建立1751天!

热门文章
文章归档
回到顶部