mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-30 23:18:20 +00:00
fd8ea4e466
If a symbolic dimension is found allow the user to provide a value, or default to 1. `python .\onnxruntime_test.py --symbolic_dims batch=1,seqlen=4 onnxruntime\test\testdata\transform\fusion\fast_gelu_use_graph_input.onnx`
118 lines
4.8 KiB
Python
118 lines
4.8 KiB
Python
#-------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
#--------------------------------------------------------------------------
|
|
|
|
import argparse
|
|
import onnxruntime as onnxrt
|
|
import numpy as np
|
|
import os
|
|
import sys
|
|
from timeit import default_timer as timer
|
|
|
|
float_dict = {'tensor(float16)': 'float16', 'tensor(float)': 'float32', 'tensor(double)': 'float64'}
|
|
|
|
integer_dict = {
|
|
'tensor(int32)': 'int32',
|
|
'tensor(int8)': 'int8',
|
|
'tensor(uint8)': 'uint8',
|
|
'tensor(int16)': 'int16',
|
|
'tensor(uint16)': 'uint16',
|
|
'tensor(int64)': 'int64',
|
|
'tensor(uint64)': 'uint64'
|
|
}
|
|
|
|
# simple test program for loading onnx model, feeding all inputs and running the model num_iters times.
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Simple ONNX Runtime Test Tool.')
|
|
parser.add_argument('model_path', help='model path')
|
|
parser.add_argument('num_iters', nargs='?', type=int, default=1000, help='model run iterations. default=1000')
|
|
parser.add_argument('--debug', action='store_true', help='pause execution to allow attaching a debugger.')
|
|
parser.add_argument('--profile', action='store_true', help='enable chrome timeline trace profiling.')
|
|
parser.add_argument('--symbolic_dims', default=None, type=lambda s: dict(x.split("=") for x in s.split(",")),
|
|
help='Comma separated name=value pairs for any symbolic dimensions in the model input. '
|
|
'e.g. --symbolic_dims batch=1,seqlen=5. '
|
|
'If not provided, the value of 1 will be used for all symbolic dimensions.')
|
|
|
|
args = parser.parse_args()
|
|
iters = args.num_iters
|
|
|
|
if args.debug:
|
|
print("Pausing execution ready for debugger to attach to pid: {}".format(os.getpid()))
|
|
print("Press key to continue.")
|
|
sys.stdin.read(1)
|
|
|
|
sess_options = None
|
|
if args.profile:
|
|
sess_options = onnxrt.SessionOptions()
|
|
sess_options.enable_profiling = True
|
|
sess_options.profile_file_prefix = os.path.basename(args.model_path)
|
|
|
|
sess = onnxrt.InferenceSession(args.model_path, sess_options)
|
|
meta = sess.get_modelmeta()
|
|
|
|
feeds = {}
|
|
for input_meta in sess.get_inputs():
|
|
# replace any symbolic dimensions
|
|
shape = []
|
|
for dim in input_meta.shape:
|
|
if not dim:
|
|
# unknown dim
|
|
shape.append(1)
|
|
elif type(dim) == str:
|
|
# symbolic dim. see if we have a value otherwise use 1
|
|
if dim in args.symbolic_dims:
|
|
shape.append(int(args.symbolic_dims[dim]))
|
|
else:
|
|
shape.append(1)
|
|
else:
|
|
shape.append(dim)
|
|
|
|
if input_meta.type in float_dict:
|
|
feeds[input_meta.name] = np.random.rand(*shape).astype(float_dict[input_meta.type])
|
|
elif input_meta.type in integer_dict:
|
|
feeds[input_meta.name] = np.random.uniform(high=1000,
|
|
size=tuple(shape)).astype(integer_dict[input_meta.type])
|
|
elif input_meta.type == 'tensor(bool)':
|
|
feeds[input_meta.name] = np.random.randint(2, size=tuple(shape)).astype('bool')
|
|
else:
|
|
print("unsupported input type {} for input {}".format(input_meta.type, input_meta.name))
|
|
sys.exit(-1)
|
|
|
|
# Starting with IR4 some initializers provide default values
|
|
# and can be overridden (available in IR4). For IR < 4 models
|
|
# the list would be empty
|
|
for initializer in sess.get_overridable_initializers():
|
|
shape = [dim if dim else 1 for dim in initializer.shape]
|
|
if initializer.type in float_dict:
|
|
feeds[initializer.name] = np.random.rand(*shape).astype(float_dict[initializer.type])
|
|
elif initializer.type in integer_dict:
|
|
feeds[initializer.name] = np.random.uniform(high=1000,
|
|
size=tuple(shape)).astype(integer_dict[initializer.type])
|
|
elif initializer.type == 'tensor(bool)':
|
|
feeds[initializer.name] = np.random.randint(2, size=tuple(shape)).astype('bool')
|
|
else:
|
|
print("unsupported initializer type {} for initializer {}".format(initializer.type, initializer.name))
|
|
sys.exit(-1)
|
|
|
|
start = timer()
|
|
for i in range(iters):
|
|
sess.run([], feeds) # fetch all outputs
|
|
end = timer()
|
|
|
|
print("model: {}".format(meta.graph_name))
|
|
print("version: {}".format(meta.version))
|
|
print("iterations: {}".format(iters))
|
|
print("avg latency: {} ms".format(((end - start) * 1000) / iters))
|
|
|
|
if args.profile:
|
|
trace_file = sess.end_profiling()
|
|
print("trace file written to: {}".format(trace_file))
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|