codekingpro/portable-devtools

1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License.
4# --------------------------------------------------------------------------
5from __future__ import annotations
6
7import argparse
8import os
9import sys
10from timeit import default_timer as timer
11
12import numpy as np
13
14import onnxruntime as onnxrt
15
16float_dict = {
17    "tensor(float16)": "float16",
18    "tensor(float)": "float32",
19    "tensor(double)": "float64",
20}
21
22integer_dict = {
23    "tensor(int32)": "int32",
24    "tensor(int8)": "int8",
25    "tensor(uint8)": "uint8",
26    "tensor(int16)": "int16",
27    "tensor(uint16)": "uint16",
28    "tensor(int64)": "int64",
29    "tensor(uint64)": "uint64",
30}
31
32
33def generate_feeds(sess, symbolic_dims: dict | None = None):
34    feeds = {}
35    symbolic_dims = symbolic_dims or {}
36    for input_meta in sess.get_inputs():
37        # replace any symbolic dimensions
38        shape = []
39        for dim in input_meta.shape:
40            if not dim:
41                # unknown dim
42                shape.append(1)
43            elif isinstance(dim, str):
44                # symbolic dim. see if we have a value otherwise use 1
45                if dim in symbolic_dims:
46                    shape.append(int(symbolic_dims[dim]))
47                else:
48                    shape.append(1)
49            else:
50                shape.append(dim)
51
52        if input_meta.type in float_dict:
53            feeds[input_meta.name] = np.random.rand(*shape).astype(float_dict[input_meta.type])
54        elif input_meta.type in integer_dict:
55            feeds[input_meta.name] = np.random.uniform(high=1000, size=tuple(shape)).astype(
56                integer_dict[input_meta.type]
57            )
58        elif input_meta.type == "tensor(bool)":
59            feeds[input_meta.name] = np.random.randint(2, size=tuple(shape)).astype("bool")
60        else:
61            print(f"unsupported input type {input_meta.type} for input {input_meta.name}")
62            sys.exit(-1)
63    return feeds
64
65
66# simple test program for loading onnx model, feeding all inputs and running the model num_iters times.
67def run_model(
68    model_path,
69    num_iters=1,
70    debug=None,
71    profile=None,
72    symbolic_dims=None,
73    feeds=None,
74    override_initializers=True,
75):
76    symbolic_dims = symbolic_dims or {}
77    if debug:
78        print(f"Pausing execution ready for debugger to attach to pid: {os.getpid()}")
79        print("Press key to continue.")
80        sys.stdin.read(1)
81
82    sess_options = None
83    if profile:
84        sess_options = onnxrt.SessionOptions()
85        sess_options.enable_profiling = True
86        sess_options.profile_file_prefix = os.path.basename(model_path)
87
88    sess = onnxrt.InferenceSession(
89        model_path,
90        sess_options=sess_options,
91        providers=onnxrt.get_available_providers(),
92    )
93    meta = sess.get_modelmeta()
94
95    if not feeds:
96        feeds = generate_feeds(sess, symbolic_dims)
97
98    if override_initializers:
99        # Starting with IR4 some initializers provide default values
100        # and can be overridden (available in IR4). For IR < 4 models
101        # the list would be empty
102        for initializer in sess.get_overridable_initializers():
103            shape = [dim if dim else 1 for dim in initializer.shape]
104            if initializer.type in float_dict:
105                feeds[initializer.name] = np.random.rand(*shape).astype(float_dict[initializer.type])
106            elif initializer.type in integer_dict:
107                feeds[initializer.name] = np.random.uniform(high=1000, size=tuple(shape)).astype(
108                    integer_dict[initializer.type]
109                )
110            elif initializer.type == "tensor(bool)":
111                feeds[initializer.name] = np.random.randint(2, size=tuple(shape)).astype("bool")
112            else:
113                print(f"unsupported initializer type {initializer.type} for initializer {initializer.name}")
114                sys.exit(-1)
115
116    start = timer()
117    for _i in range(num_iters):
118        outputs = sess.run([], feeds)  # fetch all outputs
119    end = timer()
120
121    print(f"model: {meta.graph_name}")
122    print(f"version: {meta.version}")
123    print(f"iterations: {num_iters}")
124    print(f"avg latency: {((end - start) * 1000) / num_iters} ms")
125
126    if profile:
127        trace_file = sess.end_profiling()
128        print(f"trace file written to: {trace_file}")
129
130    return 0, feeds, num_iters > 0 and outputs
131
132
133def main():
134    parser = argparse.ArgumentParser(description="Simple ONNX Runtime Test Tool.")
135    parser.add_argument("model_path", help="model path")
136    parser.add_argument(
137        "num_iters",
138        nargs="?",
139        type=int,
140        default=1000,
141        help="model run iterations. default=1000",
142    )
143    parser.add_argument(
144        "--debug",
145        action="store_true",
146        help="pause execution to allow attaching a debugger.",
147    )
148    parser.add_argument("--profile", action="store_true", help="enable chrome timeline trace profiling.")
149    parser.add_argument(
150        "--symbolic_dims",
151        default={},
152        type=lambda s: dict(x.split("=") for x in s.split(",")),
153        help="Comma separated name=value pairs for any symbolic dimensions in the model input. "
154        "e.g. --symbolic_dims batch=1,seqlen=5. "
155        "If not provided, the value of 1 will be used for all symbolic dimensions.",
156    )
157
158    args = parser.parse_args()
159    exit_code, _, _ = run_model(args.model_path, args.num_iters, args.debug, args.profile, args.symbolic_dims)
160    sys.exit(exit_code)
161
162
163if __name__ == "__main__":
164    main()
165