[RunONNXModel.py] Add options to print out input/output signatures and support big models (#2982)

* Add options to print out input/output signatures and support big models

Signed-off-by: Tung D. Le <tung@jp.ibm.com>

---------

Signed-off-by: Tung D. Le <tung@jp.ibm.com>

Author: tungld

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,24 +140,24 @@ 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",
 )
 
 parser.add_argument(
     "--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)