onnxruntime_gpu_inference_session_with_cuda_provider.py

1import onnxruntime as ort
2import numpy as np
3
4# Create a sample model (or use an existing .onnx file)
5# For this example, we assume 'model.onnx' exists. 
6# You can replace this with your own model path.
7model_path = "model.onnx"
8
9# To use GPU, you must specify the 'CUDAExecutionProvider'
10# ONNX Runtime will fall back to 'CPUExecutionProvider' if CUDA is unavailable
11providers = [
12    ('CUDAExecutionProvider', {
13        'device_id': 0,
14        'arena_extend_strategy': 'kNextPowerOfTwo',
15        'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
16        'cudnn_conv_algo_search': 'EXHAUSTIVE',
17        'do_copy_in_default_stream': True,
18    }),
19    'CPUExecutionProvider',
20]
21
22# Initialize the session
23session = ort.InferenceSession(model_path, providers=providers)
24
25# Get input names and shapes
26input_name = session.get_inputs()[0].name
27input_shape = session.get_inputs()[0].shape
28
29# Create dummy input data
30input_data = np.random.randn(*[s if isinstance(s, int) else 1 for s in input_shape]).astype(np.float32)
31
32# Run inference
33outputs = session.run(None, {input_name: input_data})
34
35print(f"Provider used: {session.get_providers()}")
36print(f"Output shape: {outputs[0].shape}")