From d35a4733def31f99c0e87d63299aa5680b1ffd02 Mon Sep 17 00:00:00 2001 From: inisis Date: Mon, 23 Dec 2024 00:12:07 +0800 Subject: [PATCH] add ort benchmark (#64) --- tests/test_benchmark.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index ad0abbc..ea2271d 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -36,12 +36,31 @@ def bench_polygraphy(input, output): return result +def bench_onnxruntime(input, output): + try: + import onnxruntime as rt + + sess_options = rt.SessionOptions() + # Set graph optimization level + sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_EXTENDED + # To enable model serialization after graph optimization set this + sess_options.optimized_model_filepath = output + rt.InferenceSession(input, sess_options) + return True + + except Exception as e: + print(e) + return None + + class TestModelZoo: def transform_and_check(self, name, filename, transformation_func, suffix, check_func): with tempfile.TemporaryDirectory() as tempdir: output_file = os.path.join(tempdir, f"{name}_{suffix}.onnx") result = transformation_func(filename, output_file) - if result.returncode == 0: + if result is None: + return None + if result or result.returncode == 0: if check_func: try: check_func(output_file) @@ -55,6 +74,7 @@ def run_model_test(self, name, filename, check_func=None): summary_list.append(self.transform_and_check(name, filename, bench_onnxslim, "onnxslim", check_func)) summary_list.append(self.transform_and_check(name, filename, bench_onnxsim, "onnxsim", check_func)) summary_list.append(self.transform_and_check(name, filename, bench_polygraphy, "polygraphy", check_func)) + summary_list.append(self.transform_and_check(name, filename, bench_onnxruntime, "onnxruntime", check_func)) summary_list = [summary for summary in summary_list if summary is not None]