Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 5 commits into from
Oct 21, 2024
Merged

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

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ced8562
Add options to print out input/output signatures and support big models
tungld Oct 18, 2024
9bef098
black format
tungld Oct 18, 2024
028ce7b
typos
tungld Oct 18, 2024
a2c8671
Merge branch 'main' into RunONNXModel_constants_to_file
tungld Oct 21, 2024
e81b5fd
Changes requested by reviewers
tungld Oct 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
217 changes: 124 additions & 93 deletions utils/RunONNXModel.py
View file
Open in desktop
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
Loading