diff --git a/utils/RunONNXModel.py b/utils/RunONNXModel.py index 017401c475..47af04f520 100755 --- a/utils/RunONNXModel.py +++ b/utils/RunONNXModel.py @@ -23,6 +23,7 @@ import json import importlib.util import shlex +import shutil from onnx import numpy_helper from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE @@ -92,6 +93,11 @@ def check_non_negative(argname, value): action="store_true", help="Print out inference outputs produced by onnx-mlir", ) +parser.add_argument( + "--print-signatures", + action="store_true", + help="Print out the input and output signatures of the model", +) parser.add_argument( "--save-onnx", metavar="PATH", @@ -134,16 +140,16 @@ def check_non_negative(argname, value): lib_group = parser.add_mutually_exclusive_group() lib_group.add_argument( - "--save-so", + "--save-model", metavar="PATH", type=str, - help="File path to save the generated shared library of" " the model", + help="Path to a folder to save the compiled model", ) lib_group.add_argument( - "--load-so", + "--load-model", metavar="PATH", type=str, - help="File path to load a generated shared library for " + help="Path to a folder to load a compiled model for " "inference, and the ONNX model will not be re-compiled", ) @@ -151,7 +157,7 @@ def check_non_negative(argname, value): "--save-ref", metavar="PATH", type=str, - help="Path to a folder to save the inputs and outputs" " in protobuf", + help="Path to a folder to save the inputs and outputs in protobuf", ) data_group = parser.add_mutually_exclusive_group() data_group.add_argument( @@ -617,17 +623,21 @@ class onnxruntime.InferenceSession(path_or_bytes: str | bytes | os.PathLike, ses Another argument, 'options' is added for onnxmlir to specify options for RunONNXModel.py """ - def __init__(self, model_name, **kwargs): + def __init__(self, model_file, **kwargs): global args if "options" in kwargs.keys(): options = kwargs["options"] args = parser.parse_args(shlex.split(options)) - if model_name: - if model_name.endswith(".onnx") or model_name.endswith(".mlir"): - args.model = model_name + if model_file: + if model_file.endswith(".onnx") or model_file.endswith(".mlir"): + args.model = model_file else: - args.load_so = compiled_name + args.load_model = compiled_name + + # Default model name that will be used for the compiled model. + # e.g. model.so, model.constants.bin, ... + self.default_model_name = "model" # Get shape information if given. # args.shape_info in the form of 'input_index:d1xd2, input_index:d1xd2' @@ -662,95 +672,111 @@ def __init__(self, model_name, **kwargs): print("Saving modified onnx model to ", args.save_onnx, "\n") onnx.save(model, args.save_onnx) - # Compile, run, and verify. - with tempfile.TemporaryDirectory() as temp_dir: - print("Temporary directory has been created at {}".format(temp_dir)) - - shared_lib_path = "" - - # If a shared library is given, use it without compiling the ONNX model. - # Otherwise, compile the ONNX model. - if args.load_so: - shared_lib_path = args.load_so - else: - print("Compiling the model ...") - # Prepare input and output paths. - output_path = os.path.join(temp_dir, "model") - shared_lib_path = os.path.join(temp_dir, "model.so") - if args.model.endswith(".onnx"): - input_model_path = os.path.join(temp_dir, "model.onnx") + # If a shared library is given, use it without compiling the ONNX model. + # Otherwise, compile the ONNX model. + if args.load_model: + self.model_dir = args.load_model + else: + # Compile the ONNX model. + self.temp_dir = tempfile.TemporaryDirectory() + print("Temporary directory has been created at {}\n".format(self.temp_dir)) + print("Compiling the model ...") + self.model_dir = self.temp_dir.name + # Prepare input and output paths. + output_path = os.path.join(self.model_dir, self.default_model_name) + if args.model.endswith(".onnx"): + if args.verify and args.verify == "onnxruntime" and args.verify_all_ops: + input_model_path = os.path.join( + self.model_dir, f"{self.default_model_name}.onnx" + ) onnx.save(model, input_model_path) - elif args.model.endswith(".mlir") or args.model.endswith(".onnxtext"): - input_model_path = args.model else: - print("Invalid input model path. Must end with .onnx or .mlir") - exit(1) - - # Prepare compiler arguments. - command_str = [ONNX_MLIR] - if args.compile_args: - command_str += args.compile_args.split() - if args.compile_using_input_shape: - # Use shapes of the reference inputs to compile the model. - assert ( - args.load_ref or args.load_ref_from_numpy - ), "No data folder given" - assert "shapeInformation" not in command_str, "shape info was set" - shape_info = "--shapeInformation=" - for i in range(len(inputs)): - shape_info += ( - str(i) - + ":" - + "x".join([str(d) for d in inputs[i].shape]) - + "," - ) - shape_info = shape_info[:-1] - command_str += [shape_info] - warning( - "the shapes of the model's inputs will be " - "changed to the shapes of the inputs in the data folder" - ) - command_str += [input_model_path] - command_str += ["-o", output_path] + input_model_path = args.model + elif args.model.endswith(".mlir") or args.model.endswith(".onnxtext"): + input_model_path = args.model + else: + print( + "Invalid input model path. Must end with .onnx or .mlir or .onnxtext" + ) + exit(1) - # Compile the model. - start = time.perf_counter() - ok, msg = execute_commands(command_str) - # Dump the compilation log into a file. - if args.log_to_file: - log_file = ( - args.log_to_file - if args.log_to_file.startswith("/") - else os.path.join(os.getcwd(), args.log_to_file) + # Prepare compiler arguments. + command_str = [ONNX_MLIR] + if args.compile_args: + command_str += args.compile_args.split() + if args.compile_using_input_shape: + # Use shapes of the reference inputs to compile the model. + assert args.load_ref or args.load_ref_from_numpy, "No data folder given" + assert "shapeInformation" not in command_str, "shape info was set" + shape_info = "--shapeInformation=" + for i in range(len(inputs)): + shape_info += ( + str(i) + ":" + "x".join([str(d) for d in inputs[i].shape]) + "," ) - print(" Compilation log is dumped into {}".format(log_file)) - with open(log_file, "w") as f: - f.write(msg) - if not ok: - print(msg) - exit(1) - end = time.perf_counter() - print(" took ", end - start, " seconds.\n") - - # Save the generated .so file of the model if required. - if args.save_so: - print("Saving the shared library to", args.save_so, "\n") - execute_commands(["rsync", "-ar", shared_lib_path, args.save_so]) - - # Exit if only compiling the model. - if args.compile_only: - exit(0) + shape_info = shape_info[:-1] + command_str += [shape_info] + warning( + "the shapes of the model's inputs will be " + "changed to the shapes of the inputs in the data folder" + ) + command_str += [input_model_path] + command_str += ["-o", output_path] - # Use the generated shared library to create an execution session. - print("Loading the compiled model ...") + # Compile the model. start = time.perf_counter() - if args.load_so: - sess = OMExecutionSession(shared_lib_path, tag="None") - else: - sess = OMExecutionSession(shared_lib_path) + ok, msg = execute_commands(command_str) + # Dump the compilation log into a file. + if args.log_to_file: + log_file = ( + args.log_to_file + if args.log_to_file.startswith("/") + else os.path.join(os.getcwd(), args.log_to_file) + ) + print(" Compilation log is dumped into {}".format(log_file)) + with open(log_file, "w") as f: + f.write(msg) + if not ok: + print(msg) + exit(1) end = time.perf_counter() print(" took ", end - start, " seconds.\n") - self.sess = sess + + # Save the generated .so and .constants.bin files of the model if required. + if args.save_model: + if not os.path.exists(args.save_model): + os.makedirs(args.save_model) + if not os.path.isdir(args.save_model): + print("Path to --save-model is not a folder") + exit(0) + shared_lib_path = self.model_dir + f"/{self.default_model_name}.so" + if os.path.exists(shared_lib_path): + print("Saving the shared library to", args.save_model) + shutil.copy2(shared_lib_path, args.save_model) + constants_file_path = os.path.join( + self.model_dir, f"{self.default_model_name}.constants.bin" + ) + if os.path.exists(constants_file_path): + print("Saving the constants file to ", args.save_model, "\n") + shutil.copy2(constants_file_path, args.save_model) + + # Exit if only compiling the model. + if args.compile_only: + exit(0) + + # Use the generated shared library to create an execution session. + start = time.perf_counter() + shared_lib_path = self.model_dir + f"/{self.default_model_name}.so" + if not os.path.exists(shared_lib_path): + print(f"Input model {shared_lib_path} does not exist") + exit(0) + print("Loading the compiled model ...") + if args.load_model: + sess = OMExecutionSession(shared_lib_path, tag="None") + else: + sess = OMExecutionSession(shared_lib_path) + end = time.perf_counter() + print(" took ", end - start, " seconds.\n") + self.sess = sess """ From onnxruntime API: @@ -788,6 +814,9 @@ def run(self, outputname, input_feed, **kwargs): output_signature = self.sess.output_signature() input_names = get_names_in_signature(input_signature) output_names = get_names_in_signature(output_signature) + if args.print_signatures: + print("Model's input signature: ", input_signature.strip()) + print("Model's output signature: ", output_signature.strip()) inputs = [] # Get input from input_feed, if input_feed is provided @@ -834,6 +863,8 @@ def run(self, outputname, input_feed, **kwargs): # Running inference. print("Running inference ...") + # Let onnx-mlir know where to find the constants file. + os.environ["OM_CONSTANT_PATH"] = self.model_dir for i in range(args.warmup): start = time.perf_counter() outs = self.sess.run(inputs) @@ -957,8 +988,8 @@ def run(self, outputname, input_feed, **kwargs): # Standalone driver def main(): - if not (args.model or args.load_so): - print("error: no input model, use argument --model and/or --load-so.") + if not (args.model or args.load_model): + print("error: no input model, use argument --model and/or --load-model.") print(parser.format_usage()) exit(1)