Skip to content

Commit

Permalink
[RunONNXModel.py] Add options to print out input/output signatures an…
Browse files Browse the repository at this point in the history 
…d support big models (#2982)

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

Signed-off-by: Tung D. Le <[email protected]>

---------

Signed-off-by: Tung D. Le <[email protected]>
  • Loading branch information
@tungld
tungld authored Oct 21, 2024
1 parent ee7feba commit c80d3e0
Showing 1 changed file with 124 additions and 93 deletions.
217 changes: 124 additions & 93 deletions utils/RunONNXModel.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit c80d3e0

Please sign in to comment.