import onnxruntime as ort
all_providers = ort.get_all_providers()
print(all_providers)
available_providers = ort.get_available_providers()
print(available_providers)
运行结果:
['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'MIGraphXExecutionProvider', 'ROCMExecutionProvider', 'OpenVINOExecutionProvider', 'DnnlExecutionProvider', 'TvmExecutionProvider', 'VitisAIExecutionProvider', 'QNNExecutionProvider', 'NnapiExecutionProvider', 'JsExecutionProvider', 'CoreMLExecutionProvider', 'ArmNNExecutionProvider', 'ACLExecutionProvider', 'DmlExecutionProvider', 'RknpuExecutionProvider', 'WebNNExecutionProvider', 'XnnpackExecutionProvider', 'CANNExecutionProvider', 'AzureExecutionProvider', 'CPUExecutionProvider']
['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
onnx通过加载模型时指定provider来选择推理设备,可以通过get_available_providers来获取可用的providers,providers是一个列表,按照先后顺序进行选择,在前一个provider无法使用的情况下会向后选择。以下是一个示例:
import os
import onnxruntime as ort
def get_available_providers():
providers = []
# 检查 CUDA 是否可用
if os.path.exists("/usr/local/cuda/bin/nvcc"):
providers.append('CUDAExecutionProvider')
# CPU 总是可用的
providers.append('CPUExecutionProvider')
return providers
def create_ort_session(model_path, providers):
# 使用指定的提供者列表创建会话
session_options = ort.SessionOptions()
session_options.enable_profiling = True # 可选,开启性能分析
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(model_path, sess_options=session_options, providers=providers)
return session
# 模型路径
model_path = 'your_model.onnx'
# 获取可用的执行提供者
providers = get_available_providers()
# 创建会话
session = create_ort_session(model_path, providers)
# 现在你可以使用 session 来运行模型了
# 例如: results = session.run(None, {'input_name': input_data})