test_onnx_model_perf.py

import numpy as np
import sys
import onnx
from onnx import numpy_helper
from onnx import mapping
from onnx import TensorProto
import onnxruntime as rt
import timeit

def create_random_data(shape, type, minvalue, maxvalue, seed):
    nptype = np.dtype(type)
    np.random.seed(seed)
    return ((maxvalue - minvalue) * np.random.sample(shape) + minvalue).astype(nptype)


if __name__ == '__main__':
    model_path = sys.argv[1]
    print("Model: ", model_path)
    num_perf_iters = 1
    if len(sys.argv) > 2:
        num_perf_iters = int(sys.argv[2])
        
    m = onnx.load(model_path)
    input_dict = {}
    output_names = []
    for i in m.graph.input:
        shape_arr = []
        shape_obj = i.type.tensor_type.shape
        for onedim in shape_obj.dim:
            if onedim.HasField("dim_param"):
                shape_arr.append(1)
            if onedim.HasField("dim_value"):
                shape_arr.append(onedim.dim_value)
        print("input name: ", i.name, " shape: ", shape_arr)
        # TODO convert type from proto to numpy type
        d = create_random_data(shape_arr, np.float32, 0, 1, None)
        input_dict[i.name] = d

    for o in m.graph.output:
        output_names.append(o.name)

    print("number of inputs: ", len(input_dict))
    print("outputs: ", output_names)

    # rt.set_default_logger_severity(0)
    sess = rt.InferenceSession(model_path)
    def run_infer():
        res = sess.run(output_names, input_dict)        
    time_taken = timeit.timeit(lambda: run_infer(), number=num_perf_iters)
    print(f"\ntime taken (secs): {time_taken/num_perf_iters:.8f}")
    print("End of program.")