From 489c03977f4865c450af37c32f08d1a634e778fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=B5=E5=B0=8F=E5=87=A1?= <2672931+whyb@users.noreply.github.com> Date: Tue, 26 Nov 2024 15:56:31 +0800 Subject: [PATCH 01/12] Added onnxslim command test. input shapes/FP32 to FP16 --- tests/test_onnxslim.py | 33 +++++++++++++++++++++++++-------- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 515f65e..d52eeb7 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -12,21 +12,38 @@ class TestFunctional: - def test_basic(self, request): - """Test the basic functionality of the slim function.""" + def __test_command_basic(self, request, in_model_path=FILENAME, out_model_name="resnet18.onnx", arg_str=""): with tempfile.TemporaryDirectory() as tempdir: - summary = summarize_model(slim(FILENAME), request.node.name) + summary = summarize_model(slim(in_model_path), request.node.name)) print_model_info_as_table(summary) - output_name = os.path.join(tempdir, "resnet18.onnx") - slim(FILENAME, output_name) - slim(FILENAME, output_name, model_check=True) - - command = f"onnxslim {FILENAME} {output_name}" + out_model_path = os.path.join(tempdir, out_model_name) + slim(in_model_path, out_model_path) + slim(in_model_path, out_model_path, model_check=True) + command = f"onnxslim {arg_str} \"{in_model_path}\" \"{out_model_path}\"" result = subprocess.run(command, shell=True, capture_output=True, text=True) output = result.stderr.strip() # Assert the expected return code print(output) assert result.returncode == 0 + return in_model_path, out_model_path + + def test_basic(self, request): + """Test the basic functionality of the slim function.""" + self.__test_command_basic(request) + + def test_input_shape_modification(self, request): + """Test the modification of input shapes.""" + input_shape_arg_str = "--input-shapes input:1,3,224,224" + self.__test_command_basic(request, FILENAME, "resnet18.onnx", input_shape_arg_str) + + def test_fp322fp16_conversion(self, request): + """Test the conversion of an ONNX model from FP32 to FP16 precision.""" + dtype_fp16_arg_str = "--dtype fp16" + _, out_model_path = self.__test_command_basic( + request, FILENAME, "resnet18_fp16.onnx", dtype_fp16_arg_str + ) + dtype_fp32_arg_str = "--dtype fp32" + self.__test_command_basic(request, out_model_path, "resnet18_fp32.onnx", dtype_fp32_arg_str) class TestFeature: From 9d8ab7001e9dad2c075e0357d2a612d0aeaa629a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=B5=E5=B0=8F=E5=87=A1?= <2672931+whyb@users.noreply.github.com> Date: Tue, 26 Nov 2024 17:46:34 +0800 Subject: [PATCH 02/12] fix build ci error --- tests/test_onnxslim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index d52eeb7..9338f83 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -14,7 +14,7 @@ class TestFunctional: def __test_command_basic(self, request, in_model_path=FILENAME, out_model_name="resnet18.onnx", arg_str=""): with tempfile.TemporaryDirectory() as tempdir: - summary = summarize_model(slim(in_model_path), request.node.name)) + summary = summarize_model(slim(in_model_path), request.node.name) print_model_info_as_table(summary) out_model_path = os.path.join(tempdir, out_model_name) slim(in_model_path, out_model_path) From 9a6efd3a9353eace1a4ad41e93309f8f03a94f04 Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Tue, 26 Nov 2024 09:47:04 +0000 Subject: [PATCH 03/12] Auto-format by https://ultralytics.com/actions --- tests/test_onnxslim.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 9338f83..5536376 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -19,7 +19,7 @@ def __test_command_basic(self, request, in_model_path=FILENAME, out_model_name=" out_model_path = os.path.join(tempdir, out_model_name) slim(in_model_path, out_model_path) slim(in_model_path, out_model_path, model_check=True) - command = f"onnxslim {arg_str} \"{in_model_path}\" \"{out_model_path}\"" + command = f'onnxslim {arg_str} "{in_model_path}" "{out_model_path}"' result = subprocess.run(command, shell=True, capture_output=True, text=True) output = result.stderr.strip() # Assert the expected return code @@ -39,9 +39,7 @@ def test_input_shape_modification(self, request): def test_fp322fp16_conversion(self, request): """Test the conversion of an ONNX model from FP32 to FP16 precision.""" dtype_fp16_arg_str = "--dtype fp16" - _, out_model_path = self.__test_command_basic( - request, FILENAME, "resnet18_fp16.onnx", dtype_fp16_arg_str - ) + _, out_model_path = self.__test_command_basic(request, FILENAME, "resnet18_fp16.onnx", dtype_fp16_arg_str) dtype_fp32_arg_str = "--dtype fp32" self.__test_command_basic(request, out_model_path, "resnet18_fp32.onnx", dtype_fp32_arg_str) From 69b872f638bdd44d9d4ff73ac732347e6bd00b60 Mon Sep 17 00:00:00 2001 From: inisis Date: Wed, 27 Nov 2024 05:36:06 +0800 Subject: [PATCH 04/12] fix argparse error --- onnxslim/argparser.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/onnxslim/argparser.py b/onnxslim/argparser.py index 95ee4fc..32ac43f 100644 --- a/onnxslim/argparser.py +++ b/onnxslim/argparser.py @@ -136,14 +136,13 @@ def _add_arguments(self): arg_type = _get_inner_type(field_def.type) default_value = field_def.default if field_def.default is not field_def.default_factory else None help_text = field_def.metadata.get("help", "") - nargs = "+" if arg_type == list else None + nargs = "+" if get_origin(arg_type) == list else None choices = field_def.metadata.get("choices", None) - if choices and default_value is not None and default_value not in choices: raise ValueError( f"Invalid default value '{default_value}' for argument '{field_name}'. Must be one of {choices}." ) - + arg_type = get_args(arg_type)[0] if get_args(arg_type) else arg_type if arg_type == bool: self.parser.add_argument( f"--{field_name.replace('_', '-')}", From b7f4605c3b44bf78eea93c5410d21716956f4f18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=B5=E5=B0=8F=E5=87=A1?= <2672931+whyb@users.noreply.github.com> Date: Wed, 27 Nov 2024 10:39:07 +0800 Subject: [PATCH 05/12] The arg_str should be placed after input_model and output_model --- tests/test_onnxslim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 5536376..7f29c35 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -19,7 +19,7 @@ def __test_command_basic(self, request, in_model_path=FILENAME, out_model_name=" out_model_path = os.path.join(tempdir, out_model_name) slim(in_model_path, out_model_path) slim(in_model_path, out_model_path, model_check=True) - command = f'onnxslim {arg_str} "{in_model_path}" "{out_model_path}"' + command = f'onnxslim "{in_model_path}" "{out_model_path}" {arg_str}' result = subprocess.run(command, shell=True, capture_output=True, text=True) output = result.stderr.strip() # Assert the expected return code From 7e53d407c5b3a6348dc8216f6c74401a46e16d58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=B5=E5=B0=8F=E5=87=A1?= <2672931+whyb@users.noreply.github.com> Date: Wed, 27 Nov 2024 11:13:44 +0800 Subject: [PATCH 06/12] change __test_command_basic function --- tests/test_onnxslim.py | 48 +++++++++++++++++++++++------------------- 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 7f29c35..52cb925 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -11,37 +11,41 @@ FILENAME = f"{MODELZOO_PATH}/resnet18/resnet18.onnx" -class TestFunctional: - def __test_command_basic(self, request, in_model_path=FILENAME, out_model_name="resnet18.onnx", arg_str=""): - with tempfile.TemporaryDirectory() as tempdir: - summary = summarize_model(slim(in_model_path), request.node.name) - print_model_info_as_table(summary) - out_model_path = os.path.join(tempdir, out_model_name) - slim(in_model_path, out_model_path) - slim(in_model_path, out_model_path, model_check=True) - command = f'onnxslim "{in_model_path}" "{out_model_path}" {arg_str}' - result = subprocess.run(command, shell=True, capture_output=True, text=True) - output = result.stderr.strip() - # Assert the expected return code - print(output) - assert result.returncode == 0 - return in_model_path, out_model_path +class TestFunctional: + def __test_command_basic(self, request, in_model_path, out_model_path, arg_str=""): + summary = summarize_model(slim(in_model_path), request.node.name) + print_model_info_as_table(summary) + slim(in_model_path, out_model_path) + slim(in_model_path, out_model_path, model_check=True) + command = f'onnxslim "{in_model_path}" "{out_model_path}" {arg_str}' + result = subprocess.run(command, shell=True, capture_output=True, text=True) + output = result.stderr.strip() + # Assert the expected return code + print(output) + assert result.returncode == 0 def test_basic(self, request): """Test the basic functionality of the slim function.""" - self.__test_command_basic(request) + with tempfile.TemporaryDirectory() as tempdir: + out_model_path = os.path.join(tempdir, "resnet18.onnx") + self.__test_command_basic(request, FILENAME, out_model_path, "") def test_input_shape_modification(self, request): """Test the modification of input shapes.""" - input_shape_arg_str = "--input-shapes input:1,3,224,224" - self.__test_command_basic(request, FILENAME, "resnet18.onnx", input_shape_arg_str) + with tempfile.TemporaryDirectory() as tempdir: + out_model_path = os.path.join(tempdir, "resnet18.onnx") + input_shape_arg_str = "--input-shapes input:1,3,224,224" + self.__test_command_basic(request, FILENAME, out_model_path, input_shape_arg_str) def test_fp322fp16_conversion(self, request): """Test the conversion of an ONNX model from FP32 to FP16 precision.""" - dtype_fp16_arg_str = "--dtype fp16" - _, out_model_path = self.__test_command_basic(request, FILENAME, "resnet18_fp16.onnx", dtype_fp16_arg_str) - dtype_fp32_arg_str = "--dtype fp32" - self.__test_command_basic(request, out_model_path, "resnet18_fp32.onnx", dtype_fp32_arg_str) + with tempfile.TemporaryDirectory() as tempdir: + out_fp16_model_path = os.path.join(tempdir, "resnet18_fp16.onnx") + out_fp32_model_path = os.path.join(tempdir, "resnet18_fp32.onnx") + dtype_fp16_arg_str = "--dtype fp16" + dtype_fp32_arg_str = "--dtype fp32" + self.__test_command_basic(request, FILENAME, out_fp16_model_path, dtype_fp16_arg_str) + self.__test_command_basic(request, out_fp16_model_path, out_fp32_model_path, dtype_fp32_arg_str) class TestFeature: From 03fadff3c9a0c8aa453dcc937be425e156518e89 Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Wed, 27 Nov 2024 03:14:12 +0000 Subject: [PATCH 07/12] Auto-format by https://ultralytics.com/actions --- tests/test_onnxslim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 52cb925..34afd1b 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -11,7 +11,7 @@ FILENAME = f"{MODELZOO_PATH}/resnet18/resnet18.onnx" -class TestFunctional: +class TestFunctional: def __test_command_basic(self, request, in_model_path, out_model_path, arg_str=""): summary = summarize_model(slim(in_model_path), request.node.name) print_model_info_as_table(summary) From 1685e1535fb07bc4e8ffea662b92bb7787a7bfa2 Mon Sep 17 00:00:00 2001 From: inisis Date: Wed, 27 Nov 2024 20:22:03 +0800 Subject: [PATCH 08/12] refactor tests --- tests/test_onnxslim.py | 142 ++++++++++++++++++++--------------------- 1 file changed, 69 insertions(+), 73 deletions(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 34afd1b..67bc405 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -3,113 +3,109 @@ import tempfile import pytest - +import numpy as np from onnxslim import slim -from onnxslim.utils import print_model_info_as_table, summarize_model +from onnxslim.utils import summarize_model MODELZOO_PATH = "/data/modelzoo" FILENAME = f"{MODELZOO_PATH}/resnet18/resnet18.onnx" class TestFunctional: - def __test_command_basic(self, request, in_model_path, out_model_path, arg_str=""): - summary = summarize_model(slim(in_model_path), request.node.name) - print_model_info_as_table(summary) - slim(in_model_path, out_model_path) - slim(in_model_path, out_model_path, model_check=True) - command = f'onnxslim "{in_model_path}" "{out_model_path}" {arg_str}' + def run_basic_test(self, in_model_path, out_model_path, **kwargs): + check_func = kwargs.get('check_func', None) + kwargs_api = kwargs.get('api', {}) + kwargs_bash = kwargs.get('bash', "") + summary = summarize_model(slim(in_model_path, **kwargs_api), in_model_path) + if check_func: check_func(summary) + + slim(in_model_path, out_model_path, **kwargs_api) + summary = summarize_model(out_model_path, out_model_path) + if check_func: check_func(summary) + + command = f'onnxslim "{in_model_path}" "{out_model_path}", {kwargs_bash}' + result = subprocess.run(command, shell=True, capture_output=True, text=True) output = result.stderr.strip() # Assert the expected return code print(output) assert result.returncode == 0 + summary = summarize_model(out_model_path, out_model_path) + if check_func: check_func(summary) + def test_basic(self, request): - """Test the basic functionality of the slim function.""" with tempfile.TemporaryDirectory() as tempdir: out_model_path = os.path.join(tempdir, "resnet18.onnx") - self.__test_command_basic(request, FILENAME, out_model_path, "") + self.run_basic_test(FILENAME, out_model_path) def test_input_shape_modification(self, request): - """Test the modification of input shapes.""" - with tempfile.TemporaryDirectory() as tempdir: - out_model_path = os.path.join(tempdir, "resnet18.onnx") - input_shape_arg_str = "--input-shapes input:1,3,224,224" - self.__test_command_basic(request, FILENAME, out_model_path, input_shape_arg_str) + def check_func(summary): + assert summary.input_info[0].shape == (1, 3, 224, 224) - def test_fp322fp16_conversion(self, request): - """Test the conversion of an ONNX model from FP32 to FP16 precision.""" with tempfile.TemporaryDirectory() as tempdir: - out_fp16_model_path = os.path.join(tempdir, "resnet18_fp16.onnx") - out_fp32_model_path = os.path.join(tempdir, "resnet18_fp32.onnx") - dtype_fp16_arg_str = "--dtype fp16" - dtype_fp32_arg_str = "--dtype fp32" - self.__test_command_basic(request, FILENAME, out_fp16_model_path, dtype_fp16_arg_str) - self.__test_command_basic(request, out_fp16_model_path, out_fp32_model_path, dtype_fp32_arg_str) - + out_model_path = os.path.join(tempdir, "resnet18.onnx") + kwargs = {} + kwargs['bash'] = "--input-shapes input:1,3,224,224" + kwargs['api'] = {'input_shapes': ["input:1,3,224,224"]} + kwargs['check_func'] = check_func + self.run_basic_test(FILENAME, out_model_path, **kwargs) -class TestFeature: - def test_input_shape_modification(self, request): - """Test the modification of input shapes.""" - summary = summarize_model(slim(FILENAME, input_shapes=["input:1,3,224,224"]), request.node.name) - print_model_info_as_table(summary) - assert summary.input_info[0].shape == (1, 3, 224, 224) + def test_input_modification(self, request): + def check_func(summary): + assert "/maxpool/MaxPool_output_0" in summary.input_maps + assert "/layer1/layer1.0/relu/Relu_output_0" in summary.input_maps with tempfile.TemporaryDirectory() as tempdir: - output_name = os.path.join(tempdir, "resnet18.onnx") - slim(FILENAME, output_name, input_shapes=["input:1,3,224,224"]) - summary = summarize_model(output_name, request.node.name) - print_model_info_as_table(summary) - assert summary.input_info[0].shape == (1, 3, 224, 224) + out_model_path = os.path.join(tempdir, "resnet18.onnx") + kwargs = {} + kwargs['bash'] = "--inputs /maxpool/MaxPool_output_0 /layer1/layer1.0/relu/Relu_output_0" + kwargs['api'] = {"inputs": ["/maxpool/MaxPool_output_0", "/layer1/layer1.0/relu/Relu_output_0"]} + kwargs['check_func'] = check_func + self.run_basic_test(FILENAME, out_model_path, **kwargs) - def test_fp162fp32_conversion(self, request): - """Test the conversion of an ONNX model from FP16 to FP32 precision.""" - import numpy as np + def test_output_modification(self, request): + def check_func(summary): + assert "/Flatten_output_0" in summary.output_maps with tempfile.TemporaryDirectory() as tempdir: - output_name = os.path.join(tempdir, "resnet18.onnx") - slim(FILENAME, output_name, input_shapes=["input:1,3,224,224"], dtype="fp16") - summary = summarize_model(output_name, request.node.name) - print_model_info_as_table(summary) + out_model_path = os.path.join(tempdir, "resnet18.onnx") + kwargs = {} + kwargs['bash'] = "--outputs /Flatten_output_0" + kwargs['api'] = {"outputs": ["/Flatten_output_0"]} + kwargs['check_func'] = check_func + self.run_basic_test(FILENAME, out_model_path, **kwargs) + + def test_dtype_conversion(self, request): + def check_func_fp16(summary): assert summary.input_info[0].dtype == np.float16 - assert summary.input_info[0].shape == (1, 3, 224, 224) - slim(output_name, output_name, dtype="fp32") - summary = summarize_model(output_name, request.node.name) - print_model_info_as_table(summary) + def check_func_fp32(summary): assert summary.input_info[0].dtype == np.float32 - assert summary.input_info[0].shape == (1, 3, 224, 224) - - def test_output_modification(self, request): - """Tests output modification.""" - summary = summarize_model(slim(FILENAME, outputs=["/Flatten_output_0"]), request.node.name) - print_model_info_as_table(summary) - assert "/Flatten_output_0" in summary.output_maps with tempfile.TemporaryDirectory() as tempdir: - output_name = os.path.join(tempdir, "resnet18.onnx") - slim(FILENAME, output_name, outputs=["/Flatten_output_0"]) - summary = summarize_model(output_name, request.node.name) - print_model_info_as_table(summary) - assert "/Flatten_output_0" in summary.output_maps + out_fp16_model_path = os.path.join(tempdir, "resnet18_fp16.onnx") + kwargs = {} + kwargs['bash'] = "--dtype fp16" + kwargs['api'] = {"dtype": "fp16"} + kwargs['check_func'] = check_func_fp16 + self.run_basic_test(FILENAME, out_fp16_model_path, **kwargs) - def test_input_modification(self, request): - """Tests input modification.""" - summary = summarize_model( - slim(FILENAME, inputs=["/maxpool/MaxPool_output_0", "/layer1/layer1.0/relu/Relu_output_0"]), - request.node.name, - ) - print_model_info_as_table(summary) - assert "/maxpool/MaxPool_output_0" in summary.input_maps - assert "/layer1/layer1.0/relu/Relu_output_0" in summary.input_maps + out_fp32_model_path = os.path.join(tempdir, "resnet18_fp32.onnx") + kwargs = {} + kwargs['bash'] = "--dtype fp32" + kwargs['api'] = {"dtype": "fp32"} + kwargs['check_func'] = check_func_fp32 + self.run_basic_test(out_fp16_model_path, out_fp32_model_path, **kwargs) + def test_save_as_external_data(self, request): with tempfile.TemporaryDirectory() as tempdir: - output_name = os.path.join(tempdir, "resnet18.onnx") - slim(FILENAME, output_name, inputs=["/maxpool/MaxPool_output_0", "/layer1/layer1.0/relu/Relu_output_0"]) - summary = summarize_model(output_name, request.node.name) - print_model_info_as_table(summary) - assert "/maxpool/MaxPool_output_0" in summary.input_maps - assert "/layer1/layer1.0/relu/Relu_output_0" in summary.input_maps + out_model_path = os.path.join(tempdir, "resnet18.onnx") + kwargs = {} + kwargs['bash'] = "--save-as-external-data" + kwargs['api'] = {"save_as_external_data": True} + self.run_basic_test(FILENAME, out_model_path, **kwargs) + assert os.path.getsize(out_model_path) < 1e5 if __name__ == "__main__": From 64aee5fd507b096b8bf24353455458a618ac55be Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Wed, 27 Nov 2024 12:23:20 +0000 Subject: [PATCH 09/12] Auto-format by https://ultralytics.com/actions --- tests/test_onnxslim.py | 52 +++++++++++++++++++++++------------------- 1 file changed, 28 insertions(+), 24 deletions(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 67bc405..3e58dbb 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -2,8 +2,9 @@ import subprocess import tempfile -import pytest import numpy as np +import pytest + from onnxslim import slim from onnxslim.utils import summarize_model @@ -13,15 +14,17 @@ class TestFunctional: def run_basic_test(self, in_model_path, out_model_path, **kwargs): - check_func = kwargs.get('check_func', None) - kwargs_api = kwargs.get('api', {}) - kwargs_bash = kwargs.get('bash', "") + check_func = kwargs.get("check_func", None) + kwargs_api = kwargs.get("api", {}) + kwargs_bash = kwargs.get("bash", "") summary = summarize_model(slim(in_model_path, **kwargs_api), in_model_path) - if check_func: check_func(summary) + if check_func: + check_func(summary) slim(in_model_path, out_model_path, **kwargs_api) summary = summarize_model(out_model_path, out_model_path) - if check_func: check_func(summary) + if check_func: + check_func(summary) command = f'onnxslim "{in_model_path}" "{out_model_path}", {kwargs_bash}' @@ -32,7 +35,8 @@ def run_basic_test(self, in_model_path, out_model_path, **kwargs): assert result.returncode == 0 summary = summarize_model(out_model_path, out_model_path) - if check_func: check_func(summary) + if check_func: + check_func(summary) def test_basic(self, request): with tempfile.TemporaryDirectory() as tempdir: @@ -46,9 +50,9 @@ def check_func(summary): with tempfile.TemporaryDirectory() as tempdir: out_model_path = os.path.join(tempdir, "resnet18.onnx") kwargs = {} - kwargs['bash'] = "--input-shapes input:1,3,224,224" - kwargs['api'] = {'input_shapes': ["input:1,3,224,224"]} - kwargs['check_func'] = check_func + kwargs["bash"] = "--input-shapes input:1,3,224,224" + kwargs["api"] = {"input_shapes": ["input:1,3,224,224"]} + kwargs["check_func"] = check_func self.run_basic_test(FILENAME, out_model_path, **kwargs) def test_input_modification(self, request): @@ -59,9 +63,9 @@ def check_func(summary): with tempfile.TemporaryDirectory() as tempdir: out_model_path = os.path.join(tempdir, "resnet18.onnx") kwargs = {} - kwargs['bash'] = "--inputs /maxpool/MaxPool_output_0 /layer1/layer1.0/relu/Relu_output_0" - kwargs['api'] = {"inputs": ["/maxpool/MaxPool_output_0", "/layer1/layer1.0/relu/Relu_output_0"]} - kwargs['check_func'] = check_func + kwargs["bash"] = "--inputs /maxpool/MaxPool_output_0 /layer1/layer1.0/relu/Relu_output_0" + kwargs["api"] = {"inputs": ["/maxpool/MaxPool_output_0", "/layer1/layer1.0/relu/Relu_output_0"]} + kwargs["check_func"] = check_func self.run_basic_test(FILENAME, out_model_path, **kwargs) def test_output_modification(self, request): @@ -71,9 +75,9 @@ def check_func(summary): with tempfile.TemporaryDirectory() as tempdir: out_model_path = os.path.join(tempdir, "resnet18.onnx") kwargs = {} - kwargs['bash'] = "--outputs /Flatten_output_0" - kwargs['api'] = {"outputs": ["/Flatten_output_0"]} - kwargs['check_func'] = check_func + kwargs["bash"] = "--outputs /Flatten_output_0" + kwargs["api"] = {"outputs": ["/Flatten_output_0"]} + kwargs["check_func"] = check_func self.run_basic_test(FILENAME, out_model_path, **kwargs) def test_dtype_conversion(self, request): @@ -86,24 +90,24 @@ def check_func_fp32(summary): with tempfile.TemporaryDirectory() as tempdir: out_fp16_model_path = os.path.join(tempdir, "resnet18_fp16.onnx") kwargs = {} - kwargs['bash'] = "--dtype fp16" - kwargs['api'] = {"dtype": "fp16"} - kwargs['check_func'] = check_func_fp16 + kwargs["bash"] = "--dtype fp16" + kwargs["api"] = {"dtype": "fp16"} + kwargs["check_func"] = check_func_fp16 self.run_basic_test(FILENAME, out_fp16_model_path, **kwargs) out_fp32_model_path = os.path.join(tempdir, "resnet18_fp32.onnx") kwargs = {} - kwargs['bash'] = "--dtype fp32" - kwargs['api'] = {"dtype": "fp32"} - kwargs['check_func'] = check_func_fp32 + kwargs["bash"] = "--dtype fp32" + kwargs["api"] = {"dtype": "fp32"} + kwargs["check_func"] = check_func_fp32 self.run_basic_test(out_fp16_model_path, out_fp32_model_path, **kwargs) def test_save_as_external_data(self, request): with tempfile.TemporaryDirectory() as tempdir: out_model_path = os.path.join(tempdir, "resnet18.onnx") kwargs = {} - kwargs['bash'] = "--save-as-external-data" - kwargs['api'] = {"save_as_external_data": True} + kwargs["bash"] = "--save-as-external-data" + kwargs["api"] = {"save_as_external_data": True} self.run_basic_test(FILENAME, out_model_path, **kwargs) assert os.path.getsize(out_model_path) < 1e5 From 961db19268b7e37f4469af7ad2bb890d7d5064a0 Mon Sep 17 00:00:00 2001 From: inisis Date: Wed, 27 Nov 2024 20:53:03 +0800 Subject: [PATCH 10/12] fix model_size bug --- onnxslim/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxslim/utils.py b/onnxslim/utils.py index 4b8c3e4..a4631f3 100644 --- a/onnxslim/utils.py +++ b/onnxslim/utils.py @@ -504,7 +504,7 @@ def save( if model_info: model_size = model.ByteSize() - model_info["model_size"] = [model_size, model_info["model_size"]] + model_info.model_size = [model_size, model_info.model_size] def check_result(raw_onnx_output, slimmed_onnx_output): From c4f181fcf239c629ca4b7cb87fbbb4d9683587d9 Mon Sep 17 00:00:00 2001 From: inisis Date: Wed, 27 Nov 2024 22:40:17 +0800 Subject: [PATCH 11/12] add more tests --- tests/test_onnxslim.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 3e58dbb..973e230 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -26,7 +26,7 @@ def run_basic_test(self, in_model_path, out_model_path, **kwargs): if check_func: check_func(summary) - command = f'onnxslim "{in_model_path}" "{out_model_path}", {kwargs_bash}' + command = f'onnxslim "{in_model_path}" "{out_model_path}" {kwargs_bash}' result = subprocess.run(command, shell=True, capture_output=True, text=True) output = result.stderr.strip() @@ -111,6 +111,24 @@ def test_save_as_external_data(self, request): self.run_basic_test(FILENAME, out_model_path, **kwargs) assert os.path.getsize(out_model_path) < 1e5 + def test_model_check(self, request): + with tempfile.TemporaryDirectory() as tempdir: + out_model_path = os.path.join(tempdir, "resnet18.onnx") + input_data = os.path.join(tempdir, "input.npy") + test_data = np.random.random((1, 3, 224 ,224)).astype(np.float32) + np.save(input_data, test_data) + kwargs = {} + kwargs["bash"] = f"--model-check --model-check-inputs input:{input_data}" + kwargs["api"] = {"model_check": True, "model_check_inputs": [f"input:{input_data}"]} + self.run_basic_test(FILENAME, out_model_path, **kwargs) + + def test_inspect(self, request): + with tempfile.TemporaryDirectory() as tempdir: + kwargs = {} + kwargs["bash"] = f"--inspect" + kwargs["api"] = {"inspect": True} + self.run_basic_test(FILENAME, FILENAME, **kwargs) + if __name__ == "__main__": pytest.main( From 815baa3eb6136420416bab7c41b3bf0d9d981585 Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Wed, 27 Nov 2024 14:40:58 +0000 Subject: [PATCH 12/12] Auto-format by https://ultralytics.com/actions --- tests/test_onnxslim.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_onnxslim.py b/tests/test_onnxslim.py index 973e230..dd7a961 100644 --- a/tests/test_onnxslim.py +++ b/tests/test_onnxslim.py @@ -115,7 +115,7 @@ def test_model_check(self, request): with tempfile.TemporaryDirectory() as tempdir: out_model_path = os.path.join(tempdir, "resnet18.onnx") input_data = os.path.join(tempdir, "input.npy") - test_data = np.random.random((1, 3, 224 ,224)).astype(np.float32) + test_data = np.random.random((1, 3, 224, 224)).astype(np.float32) np.save(input_data, test_data) kwargs = {} kwargs["bash"] = f"--model-check --model-check-inputs input:{input_data}" @@ -123,9 +123,9 @@ def test_model_check(self, request): self.run_basic_test(FILENAME, out_model_path, **kwargs) def test_inspect(self, request): - with tempfile.TemporaryDirectory() as tempdir: + with tempfile.TemporaryDirectory(): kwargs = {} - kwargs["bash"] = f"--inspect" + kwargs["bash"] = "--inspect" kwargs["api"] = {"inspect": True} self.run_basic_test(FILENAME, FILENAME, **kwargs)