ONNX Set Device ID with CUDA
作者:XD / 发表: 2023年2月15日 08:04 / 更新: 2023年2月15日 08:04 / 编程笔记 / 阅读量:1788
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)