add disable_onnx_optimisation flag and address PR comments

luxonis · Dec 12, 2024 · 887ed47 · 887ed47
1 parent 7237a6c
commit 887ed47
Show file tree

Hide file tree

Showing 6 changed files with 179 additions and 168 deletions.
diff --git a/modelconverter/cli/utils.py b/modelconverter/cli/utils.py
@@ -381,6 +381,7 @@ def get_target_specific_options(
     json_cfg = cfg.model_dump(mode="json")
     options = {
         "disable_onnx_simplification": cfg.disable_onnx_simplification,
+        "disable_onnx_optimisation": cfg.disable_onnx_optimisation,
         "inputs": json_cfg["inputs"],
     }
     if target == "rvc4":

diff --git a/modelconverter/packages/base_exporter.py b/modelconverter/packages/base_exporter.py
@@ -39,6 +39,7 @@ def __init__(
         self.outputs = {out.name: out for out in config.outputs}
         self.keep_intermediate_outputs = config.keep_intermediate_outputs
         self.disable_onnx_simplification = config.disable_onnx_simplification
+        self.disable_onnx_optimisation = config.disable_onnx_optimisation
 
         self.model_name = self.input_model.stem
 

diff --git a/modelconverter/packages/rvc4/exporter.py b/modelconverter/packages/rvc4/exporter.py
@@ -60,18 +60,22 @@ def __init__(self, config: SingleStageConfig, output_dir: Path):
                 self.inputs,
             )
 
-            onnx_modifier = ONNXModifier(
-                model_path=self.input_model,
-                output_path=self._attach_suffix(
-                    self.input_model, "modified_optimised.onnx"
-                ),
-            )
-            onnx_modifier.modify_onnx()
-            if onnx_modifier.compare_outputs():
-                logger.info("ONNX model has been optimised for RVC4.")
-                shutil.move(onnx_modifier.output_path, self.input_model)
-            else:
-                os.remove(onnx_modifier.output_path)
+            if not config.disable_onnx_optimisation:
+                onnx_modifier = ONNXModifier(
+                    model_path=self.input_model,
+                    output_path=self._attach_suffix(
+                        self.input_model, "modified_optimised.onnx"
+                    ),
+                )
+
+                if (
+                    onnx_modifier.modify_onnx()
+                    and onnx_modifier.compare_outputs()
+                ):
+                    logger.info("ONNX model has been optimised for RVC4.")
+                    shutil.move(onnx_modifier.output_path, self.input_model)
+                else:
+                    os.remove(onnx_modifier.output_path)
         else:
             logger.warning(
                 "Input file type is not ONNX. Skipping pre-processing."

diff --git a/modelconverter/utils/config.py b/modelconverter/utils/config.py
@@ -280,6 +280,7 @@ class SingleStageConfig(CustomBaseModel):
 
     keep_intermediate_outputs: bool = True
     disable_onnx_simplification: bool = False
+    disable_onnx_optimisation: bool = False
     output_remote_url: Optional[str] = None
     put_file_plugin: Optional[str] = None