Merge pull request #2 from PINTO0309/feat/disable_onnxsim

`--disable_onnxsim` Suppress the execution of onnxsim on the backend and dare to leave redundant processing

README.md

@@ -40,6 +40,7 @@ usage:
   -if INPUT_ONNX_FILE_PATH
   -of OUTPUT_ONNX_FILE_PATH
   -ics INITIALIZATION_CHARACTER_STRING
+  [-dos]
   [-n]
 
 optional arguments:
@@ -56,6 +57,9 @@ optional arguments:
       String to initialize batch size. "-1" or "N" or "xxx", etc...
       Default: '-1'
 
+  -dos, --disable_onnxsim
+      Suppress the execution of onnxsim on the backend and dare to leave redundant processing.
+
   -n, --non_verbose
       Do not show all information logs. Only error logs are displayed.
 ```
@@ -72,7 +76,8 @@ initialize(
   onnx_graph: Union[onnx.onnx_ml_pb2.ModelProto, NoneType] = None,
   output_onnx_file_path: Union[str, NoneType] = '',
   initialization_character_string: Union[str, NoneType] = '-1',
-  non_verbose: Union[bool, NoneType] = False
+  non_verbose: Union[bool, NoneType] = False,
+  disable_onnxsim: Union[bool, NoneType] = False,
 ) -> onnx.onnx_ml_pb2.ModelProto
 
     Parameters
@@ -95,6 +100,10 @@ initialize(
         String to initialize batch size. "-1" or "N" or "xxx", etc...
         Default: '-1'
 
+    disable_onnxsim: Optional[bool]
+        Suppress the execution of onnxsim on the backend and dare to leave redundant processing.
+        Default: False
+
     non_verbose: Optional[bool]
         Do not show all information logs. Only error logs are displayed.
         Default: False

sbi4onnx/__init__.py

@@ -1,3 +1,3 @@
 from sbi4onnx.onnx_batchsize_initialize import initialize, main
 
-__version__ = '1.0.4'
+__version__ = '1.0.5'

sbi4onnx/onnx_batchsize_initialize.py

@@ -39,6 +39,7 @@ def initialize(
     onnx_graph: Optional[onnx.ModelProto] = None,
     output_onnx_file_path: Optional[str] = '',
     initialization_character_string: Optional[str] = '-1',
+    disable_onnxsim: Optional[bool] = False,
     non_verbose: Optional[bool] = False,
 ) -> onnx.ModelProto:
     """
@@ -62,6 +63,10 @@ def initialize(
         String to initialize batch size. "-1" or "N" or "xxx", etc...\n
         Default: '-1'
 
+    disable_onnxsim: Optional[bool]
+        Suppress the execution of onnxsim on the backend and dare to leave redundant processing.\n\
+        Default: False
+
     non_verbose: Optional[bool]
         Do not show all information logs. Only error logs are displayed.\n\
         Default: False
@@ -92,7 +97,10 @@ def initialize(
     if not onnx_graph:
         onnx_graph = onnx.load(input_onnx_file_path)
     try:
-        onnx_graph, _ = simplify(onnx_graph)
+        # onnx-simplifier does not support optimization of ONNX files containing custom domains,
+        # so skip simplify if it contains custom domains
+        if not disable_onnxsim:
+            onnx_graph, _ = simplify(onnx_graph)
     except:
         pass
     graph = gs.import_onnx(onnx_graph)
@@ -169,6 +177,12 @@ def main():
             'String to initialize batch size. "-1" or "N" or "xxx", etc... \n'+
             'Default: \'-1\''
     )
+    parser.add_argument(
+        '-dos',
+        '--disable_onnxsim',
+        action='store_true',
+        help='Suppress the execution of onnxsim on the backend and dare to leave redundant processing.'
+    )
     parser.add_argument(
         '-n',
         '--non_verbose',
@@ -180,6 +194,7 @@ def main():
     input_onnx_file_path = args.input_onnx_file_path
     output_onnx_file_path = args.output_onnx_file_path
     initialization_character_string = args.initialization_character_string
+    disable_onnxsim = args.disable_onnxsim
     non_verbose = args.non_verbose
 
     # Load
@@ -191,6 +206,7 @@ def main():
         onnx_graph=onnx_graph,
         output_onnx_file_path=output_onnx_file_path,
         initialization_character_string=initialization_character_string,
+        disable_onnxsim=disable_onnxsim,
         non_verbose=non_verbose,
     )