ONNX Runtime

import onnxruntime as ort
import numpy as np
import os

# NOTE: This example assumes you have an ONNX model file named 'model.onnx'.
# You can typically export models from frameworks like PyTorch or TensorFlow to ONNX format.
# For a runnable example, you'd need to create a dummy model:
# e.g., using ONNX library:
# import onnx
# from onnx import TensorProto
# from onnx.helper import make_model, make_node, make_graph, make_tensor_value_info
# X = make_tensor_value_info('input', TensorProto.FLOAT, [None, 2])
# Y = make_tensor_value_info('output', TensorProto.FLOAT, [None, 2])
# node = make_node('Add', ['input', 'input'], ['output'])
# graph = make_graph([node], 'simple-graph', [X], [Y])
# onnx_model = make_model(graph)
# onnx.save(onnx_model, 'model.onnx')

model_path = os.environ.get('ONNX_MODEL_PATH', 'model.onnx')

# 1. Create an InferenceSession
# For GPU, add providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
session = ort.InferenceSession(model_path, providers=ort.get_available_providers())

print("Model inputs:")
for input_meta in session.get_inputs():
    print(f"  Name: {input_meta.name}, Shape: {input_meta.shape}, Type: {input_meta.type}")

print("\nModel outputs:")
for output_meta in session.get_outputs():
    print(f"  Name: {output_meta.name}, Shape: {output_meta.shape}, Type: {output_meta.type}")

# 2. Prepare input data (example for a model expecting a float32 array)
# Assuming the first input expects a 2D float32 array, e.g., shape (1, 2)
input_name = session.get_inputs()[0].name
input_shape = [dim if isinstance(dim, int) else 1 for dim in session.get_inputs()[0].shape] # Handle dynamic shapes
input_data = np.random.randn(*input_shape).astype(np.float32)

# 3. Run inference
outputs = session.run(None, {input_name: input_data})

# 4. Process outputs
print(f"\nOutput data type: {outputs[0].dtype}")
print(f"Output shape: {outputs[0].shape}")
print(f"Output data (first 5 elements): {outputs[0].flatten()[:5]}")