Run ONNX model on different devices

onnx_device.py

import onnxruntime
import numpy as np
import time
from onnxruntime.datasets import get_example

# ONNXRuntime API Documentation
# https://microsoft.github.io/onnxruntime/python/api_summary.html

log_severity_level = 0
model_name = "logreg_iris.onnx"

sess_options = onnxruntime.SessionOptions()
# Log severity level for a particular Run() invocation. 0:Verbose, 1:Info, 2:Warning. 3:Error, 4:Fatal. Default is 2.
sess_options.log_severity_level=log_severity_level

# Available official model examples
# https://github.com/microsoft/onnxruntime/tree/master/onnxruntime/python/datasets
example_model = get_example(model_name)
sess = onnxruntime.InferenceSession(example_model, sess_options=sess_options)

input_name = sess.get_inputs()[0].name
input_shape = sess.get_inputs()[0].shape
input_type = sess.get_inputs()[0].type

print(input_name, input_shape, input_type)

output_name = sess.get_outputs()[0].name
output_shape = sess.get_outputs()[0].shape
output_type = sess.get_outputs()[0].type

print(output_name, output_shape, output_type)

x = np.random.random(input_shape)
x = x.astype(np.float32)

print(sess.get_providers())

sess.set_providers(['CPUExecutionProvider'])
start_time = time.time()
for i in range(1000):
    result = sess.run([output_name], {input_name: x})
end_time = time.time()
time_elapsed = end_time - start_time
print(time_elapsed)

sess.set_providers(['CUDAExecutionProvider'])
start_time = time.time()
for i in range(1000):
    result = sess.run([output_name], {input_name: x})
end_time = time.time()
time_elapsed = end_time - start_time
print(time_elapsed)