Option to Specify --compress_to_fp16 (#19)

luxonis · Jul 25, 2024 · 6e157e7 · 6e157e7
1 parent f6542e3
commit 6e157e7
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 0 deletions.
diff --git a/modelconverter/packages/rvc2/exporter.py b/modelconverter/packages/rvc2/exporter.py
@@ -33,6 +33,7 @@ class RVC2Exporter(Exporter):
 
     def __init__(self, config: SingleStageConfig, output_dir: Path):
         super().__init__(config=config, output_dir=output_dir)
+        self.compress_to_fp16 = config.rvc2.compress_to_fp16
         self.number_of_shaves = config.rvc2.number_of_shaves
         self.number_of_cmx_slices = config.rvc2.number_of_cmx_slices
         self.superblob = config.rvc2.superblob
@@ -54,6 +55,8 @@ def _export_openvino_ir(self) -> Path:
         self._add_args(
             args, ["--output", ",".join(name for name in self.outputs)]
         )
+        if self.compress_to_fp16:
+            self._add_args(args, ["--compress_to_fp16"])
 
         if "--input" not in args:
             inp_str = ""

diff --git a/modelconverter/utils/config.py b/modelconverter/utils/config.py
@@ -260,6 +260,7 @@ class HailoConfig(TargetConfig):
 class BlobBaseConfig(TargetConfig):
     mo_args: List[str] = []
     compile_tool_args: List[str] = []
+    compress_to_fp16: bool = True
 
 
 class RVC2Config(BlobBaseConfig):

diff --git a/shared_with_container/configs/defaults.yaml b/shared_with_container/configs/defaults.yaml
@@ -150,6 +150,9 @@ stages:
       # Produces .superblob file instead of regular .blob.
       superblob: true
 
+      # If the original model has FP32 weights or biases, they are compressed to FP16. All intermediate data is kept in original precision.
+      compress_to_fp16: true
+
     # --- RVC3-Specific Argument ---
     rvc3:
       # Target device for POT. Can be one of { VPU, ANY }
@@ -163,6 +166,9 @@ stages:
       # List of additional arguments to pass to the compile_tool.
       compile_tool_args: []
 
+      # If the original model has FP32 weights or biases, they are compressed to FP16. All intermediate data is kept in original precision.
+      compress_to_fp16: true
+
     # --- RVC4-Specific Arguments ---
     rvc4:
       # List of additional arguments to pass to SNPE onnx-to-dlc.

diff --git a/tests/test_utils/test_config.py b/tests/test_utils/test_config.py
@@ -35,12 +35,14 @@
         "number_of_shaves": 8,
         "number_of_cmx_slices": 8,
         "disable_calibration": False,
+        "compress_to_fp16": True,
     },
     "rvc3": {
         "mo_args": [],
         "compile_tool_args": [],
         "pot_target_device": PotDevice.VPU,
         "disable_calibration": False,
+        "compress_to_fp16": True,
     },
     "rvc4": {
         "snpe_onnx_to_dlc_args": [],
@@ -249,6 +251,7 @@ def test_correct():
                 "compile_tool_args": [],
                 "pot_target_device": PotDevice.VPU,
                 "disable_calibration": False,
+                "compress_to_fp16": True,
             },
             "rvc4": {**DEFAULT_TARGET_CONFIGS["rvc4"]},
             "hailo": {