onnxruntime

1
"""
2
Example python code for creating a model with a single operator and performance testing it with various
3
input combinations.
4
"""
5

6
import time
7
import timeit
8

9
import numpy as np
10
import onnx
11

12
# if you copy this script elsewhere you may need to add the tools\python dir to the sys.path for this
13
# import to work.
14
# e.g. sys.path.append(r'<path to onnxruntime source>\tools\python')
15
import ort_test_dir_utils
16
from onnx import TensorProto, helper
17

18
import onnxruntime as rt
19

20
# make input deterministic
21
np.random.seed(123)
22

23

24
#
25
# Example code to create a model with just the operator to test. Adjust as necessary for what you want to test.
26
#
27
def create_model(model_name):
28
    graph_def = helper.make_graph(
29
        nodes=[
30
            helper.make_node(
31
                op_type="TopK",
32
                inputs=["X", "K"],
33
                outputs=["Values", "Indices"],
34
                name="topk",
35
                # attributes are also key-value pairs using the attribute name and appropriate type
36
                largest=1,
37
            ),
38
        ],
39
        name="test-model",
40
        inputs=[
41
            # create inputs with symbolic dims so we can use any input sizes
42
            helper.make_tensor_value_info("X", TensorProto.FLOAT, ["batch", "items"]),
43
            helper.make_tensor_value_info("K", TensorProto.INT64, [1]),
44
        ],
45
        outputs=[
46
            helper.make_tensor_value_info("Values", TensorProto.FLOAT, ["batch", "k"]),
47
            helper.make_tensor_value_info("Indices", TensorProto.INT64, ["batch", "k"]),
48
        ],
49
        initializer=[],
50
    )
51

52
    model = helper.make_model(graph_def, opset_imports=[helper.make_operatorsetid("", 11)])
53
    onnx.checker.check_model(model)
54

55
    onnx.save_model(model, model_name)
56

57

58
#
59
# Example code to create random input. Adjust as necessary for the input your model requires
60
#
61
def create_test_input(n, num_items, k):
62
    x = np.random.randn(n, num_items).astype(np.float32)
63
    k_in = np.asarray([k]).astype(np.int64)
64
    inputs = {"X": x, "K": k_in}
65

66
    return inputs
67

68

69
#
70
# Example code that tests various combinations of input sizes.
71
#
72
def run_perf_tests(model_path, num_threads=1):
73
    so = rt.SessionOptions()
74
    so.intra_op_num_threads = num_threads
75
    sess = rt.InferenceSession(model_path, sess_options=so)
76

77
    batches = [10, 25, 50]
78
    batch_size = [8, 16, 32, 64, 128, 256, 512, 1024, 2048]
79
    k_vals = [1, 2, 4, 6, 8, 16, 24, 32, 48, 64, 128]
80

81
    # exploit scope to access variables from below for each iteration
82
    def run_test():
83
        num_seconds = 1 * 1000 * 1000 * 1000  # seconds in ns
84
        iters = 0
85
        total = 0
86
        total_iters = 0
87

88
        # For a simple model execution can be faster than time.time_ns() updates. Due to this we want to estimate
89
        # a number of iterations per measurement.
90
        # Estimate based on iterations in 5ms, but note that 5ms includes all the time_ns calls
91
        # which are excluded in the real measurement. The actual time that many iterations
92
        # takes will be much lower if the individual execution time is very small.
93
        start = time.time_ns()
94
        while time.time_ns() - start < 5 * 1000 * 1000:  # 5 ms
95
            sess.run(None, inputs)
96
            iters += 1
97

98
        # run the model and measure time after 'iters' calls
99
        while total < num_seconds:
100
            start = time.time_ns()
101
            for _i in range(iters):
102
                # ignore the outputs as we're not validating them in a performance test
103
                sess.run(None, inputs)
104
            end = time.time_ns()
105
            assert end - start > 0
106
            total += end - start
107
            total_iters += iters
108

109
        # Adjust the output you want as needed
110
        print(f"n={n},items={num_items},k={k},avg:{total / total_iters:.4f}")
111

112
    # combine the various input parameters and create input for each test
113
    for n in batches:
114
        for num_items in batch_size:
115
            for k in k_vals:
116
                if k < num_items:
117
                    # adjust as necessary for the inputs your model requires
118
                    inputs = create_test_input(n, num_items, k)
119

120
                    # use timeit to disable gc etc. but let each test measure total time and average time
121
                    # as multiple iterations may be required between each measurement
122
                    timeit.timeit(lambda: run_test(), number=1)
123

124

125
#
126
# example for creating a test directory for use with onnx_test_runner or onnxruntime_perf_test
127
# so that the model can be easily run directly or from a debugger.
128
#
129
def create_example_test_directory():
130
    # fill in the inputs that we want to use specific values for
131
    input_data = {}
132
    input_data["K"] = np.asarray([64]).astype(np.int64)
133

134
    # provide symbolic dim values as needed
135
    symbolic_dim_values = {"batch": 25, "items": 256}
136

137
    # create the directory. random input will be created for any missing inputs.
138
    # the model will be run and the output will be saved as expected output for future runs
139
    ort_test_dir_utils.create_test_dir("topk.onnx", "PerfTests", "test1", input_data, symbolic_dim_values)
140

141

142
# this will create the model file in the current directory
143
create_model("topk.onnx")
144

145
# this will create a test directory that can be used with onnx_test_runner or onnxruntime_perf_test
146
create_example_test_directory()
147

148
# this can loop over various combinations of input, using the specified number of threads
149
run_perf_tests("topk.onnx", 1)
150