onnxruntime

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]}")