From 7a3d34f4f123a2b9384a91cf6545e17e2e0bca3f Mon Sep 17 00:00:00 2001 From: inisis Date: Sun, 22 Dec 2024 01:51:00 +0800 Subject: [PATCH] add ort benchmark --- tests/test_benchmark.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index ad0abbc..daf5f9a 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -35,13 +35,30 @@ def bench_polygraphy(input, output): result = bench_main(command) 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 + session = 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 +72,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]