From ced856210a4b039efa0de4017b3ef9fc2c1bd175 Mon Sep 17 00:00:00 2001 From: "Tung D. Le" Date: Fri, 18 Oct 2024 04:10:49 -0400 Subject: [PATCH 1/4] Add options to print out input/output signatures and support big models Signed-off-by: Tung D. Le --- utils/RunONNXModel.py | 229 +++++++++++++++++++++++++----------------- 1 file changed, 136 insertions(+), 93 deletions(-) diff --git a/utils/RunONNXModel.py b/utils/RunONNXModel.py index 017401c475..3c4081e814 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 @@ -87,11 +88,21 @@ def check_non_negative(argname, value): " the inputs in the reference folder set by --load-ref", ) parser.add_argument("--print-input", action="store_true", help="Print out inputs") +parser.add_argument( + "--print-input-signature", + action="store_true", + help="Print out input signature of the model", +) parser.add_argument( "--print-output", action="store_true", help="Print out inference outputs produced by onnx-mlir", ) +parser.add_argument( + "--print-output-signature", + action="store_true", + help="Print out output signature of the model", +) parser.add_argument( "--save-onnx", metavar="PATH", @@ -134,16 +145,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 +162,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 +628,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 +677,117 @@ 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_path = 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_path = self.temp_dir.name + # Prepare input and output paths. + output_path = os.path.join(self.model_path, 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_path, 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) + shared_lib_path = self.model_path + 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_path, 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_path + f"/{self.default_model_name}.so" + if not os.path.exists(shared_lib_path): + print(f"Input model {shared_lib_path} does not exists") + 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 +825,10 @@ 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_input_signature: + print("Model's input signature: ", input_signature) + if args.print_output_signature: + print("Model's output signature: ", output_signature) inputs = [] # Get input from input_feed, if input_feed is provided @@ -834,6 +875,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_path for i in range(args.warmup): start = time.perf_counter() outs = self.sess.run(inputs) @@ -957,8 +1000,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) From 9bef098ecbca83f76584f385c9cce6e37978a090 Mon Sep 17 00:00:00 2001 From: "Tung D. Le" Date: Fri, 18 Oct 2024 04:14:54 -0400 Subject: [PATCH 2/4] black format Signed-off-by: Tung D. Le --- utils/RunONNXModel.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/utils/RunONNXModel.py b/utils/RunONNXModel.py index 3c4081e814..89be280408 100755 --- a/utils/RunONNXModel.py +++ b/utils/RunONNXModel.py @@ -690,11 +690,7 @@ def __init__(self, model_file, **kwargs): # Prepare input and output paths. output_path = os.path.join(self.model_path, self.default_model_name) if args.model.endswith(".onnx"): - if ( - args.verify - and args.verify == "onnxruntime" - and args.verify_all_ops - ): + if args.verify and args.verify == "onnxruntime" and args.verify_all_ops: input_model_path = os.path.join( self.model_path, f"{self.default_model_name}.onnx" ) @@ -715,17 +711,12 @@ def __init__(self, model_file, **kwargs): 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 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]) - + "," + str(i) + ":" + "x".join([str(d) for d in inputs[i].shape]) + "," ) shape_info = shape_info[:-1] command_str += [shape_info] From 028ce7b68e08ef6c1d057de8cd906b111419385f Mon Sep 17 00:00:00 2001 From: "Tung D. Le" Date: Fri, 18 Oct 2024 04:17:26 -0400 Subject: [PATCH 3/4] typos Signed-off-by: Tung D. Le --- utils/RunONNXModel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/RunONNXModel.py b/utils/RunONNXModel.py index 89be280408..523f3eb493 100755 --- a/utils/RunONNXModel.py +++ b/utils/RunONNXModel.py @@ -769,7 +769,7 @@ def __init__(self, model_file, **kwargs): start = time.perf_counter() shared_lib_path = self.model_path + f"/{self.default_model_name}.so" if not os.path.exists(shared_lib_path): - print(f"Input model {shared_lib_path} does not exists") + print(f"Input model {shared_lib_path} does not exist") exit(0) print("Loading the compiled model ...") if args.load_model: From e81b5fd1c242b2bc400dce9dbce7f94d976defb5 Mon Sep 17 00:00:00 2001 From: "Tung D. Le" Date: Mon, 21 Oct 2024 01:15:28 -0400 Subject: [PATCH 4/4] Changes requested by reviewers Signed-off-by: Tung D. Le --- utils/RunONNXModel.py | 35 ++++++++++++++++------------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/utils/RunONNXModel.py b/utils/RunONNXModel.py index 523f3eb493..47af04f520 100755 --- a/utils/RunONNXModel.py +++ b/utils/RunONNXModel.py @@ -88,20 +88,15 @@ def check_non_negative(argname, value): " the inputs in the reference folder set by --load-ref", ) parser.add_argument("--print-input", action="store_true", help="Print out inputs") -parser.add_argument( - "--print-input-signature", - action="store_true", - help="Print out input signature of the model", -) parser.add_argument( "--print-output", action="store_true", help="Print out inference outputs produced by onnx-mlir", ) parser.add_argument( - "--print-output-signature", + "--print-signatures", action="store_true", - help="Print out output signature of the model", + help="Print out the input and output signatures of the model", ) parser.add_argument( "--save-onnx", @@ -680,19 +675,19 @@ def __init__(self, model_file, **kwargs): # If a shared library is given, use it without compiling the ONNX model. # Otherwise, compile the ONNX model. if args.load_model: - self.model_path = 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_path = self.temp_dir.name + self.model_dir = self.temp_dir.name # Prepare input and output paths. - output_path = os.path.join(self.model_path, self.default_model_name) + 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_path, f"{self.default_model_name}.onnx" + self.model_dir, f"{self.default_model_name}.onnx" ) onnx.save(model, input_model_path) else: @@ -750,12 +745,15 @@ def __init__(self, model_file, **kwargs): if args.save_model: if not os.path.exists(args.save_model): os.makedirs(args.save_model) - shared_lib_path = self.model_path + f"/{self.default_model_name}.so" + 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_path, f"{self.default_model_name}.constants.bin" + 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") @@ -767,7 +765,7 @@ def __init__(self, model_file, **kwargs): # Use the generated shared library to create an execution session. start = time.perf_counter() - shared_lib_path = self.model_path + f"/{self.default_model_name}.so" + 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) @@ -816,10 +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_input_signature: - print("Model's input signature: ", input_signature) - if args.print_output_signature: - print("Model's output 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 @@ -867,7 +864,7 @@ 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_path + os.environ["OM_CONSTANT_PATH"] = self.model_dir for i in range(args.warmup): start = time.perf_counter() outs = self.sess.run(inputs)