[Enhancement] Additional arguments support for OpenVINO Model Optimizer (#178)

* Add mo args.

* [Docs]: update docs and argument descriptions (#196)

* bump version to v0.4.0

* update docs and argument descriptions

* revert version change

* fix unnecessary change of config for dynamic exportation (#199)

* fix mmcls get classes (#215)

* fix mmcls get classes

* resolve comment

* resolve comment

* Add ModelOptimizerOptions.

* Fix merge bugs.

* Update mmpose.md (#224)

* [Dostring]add example in apis docstring (#214)

* add example in apis docstring

* add backend example in docstring

* rm blank line

* Fixed get_mo_options_from_cfg args

* fix l2norm test

Co-authored-by: RunningLeon <[email protected]>
Co-authored-by: Haofan Wang <[email protected]>
Co-authored-by: VVsssssk <[email protected]>
Co-authored-by: grimoire <[email protected]>

Author: 5 people

docs/en/backends/openvino.md

@@ -63,6 +63,28 @@ Notes:
 the RoiAlign operation is replaced with the [ExperimentalDetectronROIFeatureExtractor](https://docs.openvinotoolkit.org/latest/openvino_docs_ops_detection_ExperimentalDetectronROIFeatureExtractor_6.html) operation in the ONNX graph.
 - Models "VFNet" and "Faster R-CNN + DCN" use the custom "DeformableConv2D" operation.
 
+### Deployment config
+
+With the deployment config, you can specify additional options for the Model Optimizer.
+To do this, add the necessary parameters to the `backend_config.mo_options` in the fields `args` (for parameters with values) and `flags` (for flags).
+
+Example:
+```python
+backend_config = dict(
+    mo_options=dict(
+        args=dict({
+            '--mean_values': [0, 0, 0],
+            '--scale_values': [255, 255, 255],
+            '--data_type': 'FP32',
+        }),
+        flags=['--disable_fusing'],
+    )
+)
+```
+
+Information about the possible parameters for the Model Optimizer can be found in the [documentation](https://docs.openvino.ai/latest/openvino_docs_MO_DG_prepare_model_convert_model_Converting_Model.html).
+
+
 ### FAQs
 
 - None

mmdeploy/apis/openvino/__init__.py

@@ -6,7 +6,8 @@
 if is_available():
     from mmdeploy.backend.openvino.onnx2openvino import (get_output_model_file,
                                                          onnx2openvino)
-    from .utils import get_input_info_from_cfg
+    from .utils import get_input_info_from_cfg, get_mo_options_from_cfg
     __all__ += [
-        'onnx2openvino', 'get_output_model_file', 'get_input_info_from_cfg'
+        'onnx2openvino', 'get_output_model_file', 'get_input_info_from_cfg',
+        'get_mo_options_from_cfg'
     ]

mmdeploy/apis/openvino/utils.py

@@ -3,8 +3,9 @@
 
 import mmcv
 
+from mmdeploy.backend.openvino import ModelOptimizerOptions
 from mmdeploy.utils import get_model_inputs
-from mmdeploy.utils.config_utils import get_ir_config
+from mmdeploy.utils.config_utils import get_backend_config, get_ir_config
 
 
 def update_input_names(input_info: Dict[str, List],
@@ -50,3 +51,19 @@ def get_input_info_from_cfg(deploy_cfg: mmcv.Config) -> Dict[str, List]:
                 input_info = dict(zip(input_names, input_info))
             input_info = update_input_names(input_info, input_names)
     return input_info
+
+
+def get_mo_options_from_cfg(deploy_cfg: mmcv.Config) -> ModelOptimizerOptions:
+    """Get additional parameters for the Model Optimizer from the deploy
+    config.
+
+    Args:
+        deploy_cfg (mmcv.Config): Deployment config.
+
+    Returns:
+        ModelOptimizerOptions: A class that will contain additional arguments.
+    """
+    backend_config = get_backend_config(deploy_cfg)
+    mo_options = backend_config.get('mo_options', None)
+    mo_options = ModelOptimizerOptions(mo_options)
+    return mo_options

mmdeploy/backend/openvino/__init__.py

@@ -13,5 +13,8 @@ def is_available() -> bool:
 
 if is_available():
     from .onnx2openvino import get_output_model_file
+    from .utils import ModelOptimizerOptions
     from .wrapper import OpenVINOWrapper
-    __all__ = ['OpenVINOWrapper', 'get_output_model_file']
+    __all__ = [
+        'OpenVINOWrapper', 'get_output_model_file', 'ModelOptimizerOptions'
+    ]

mmdeploy/backend/openvino/onnx2openvino.py

@@ -2,12 +2,13 @@
 import os.path as osp
 import subprocess
 from subprocess import PIPE, CalledProcessError, run
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 import mmcv
 import torch
 
 from mmdeploy.utils import get_root_logger
+from .utils import ModelOptimizerOptions
 
 
 def get_mo_command() -> str:
@@ -55,7 +56,10 @@ def get_output_model_file(onnx_path: str, work_dir: str) -> str:
 
 
 def onnx2openvino(input_info: Dict[str, Union[List[int], torch.Size]],
-                  output_names: List[str], onnx_path: str, work_dir: str):
+                  output_names: List[str],
+                  onnx_path: str,
+                  work_dir: str,
+                  mo_options: Optional[ModelOptimizerOptions] = None):
     """Convert ONNX to OpenVINO.
 
     Examples:
@@ -72,8 +76,9 @@ def onnx2openvino(input_info: Dict[str, Union[List[int], torch.Size]],
         output_names (List[str]): Output names. Example: ['dets', 'labels'].
         onnx_path (str): The path to the onnx model.
         work_dir (str): The path to the directory for saving the results.
+        mo_options (None | ModelOptimizerOptions): The class with
+            additional arguments for the Model Optimizer.
     """
-
     input_names = ','.join(input_info.keys())
     input_shapes = ','.join(str(list(elem)) for elem in input_info.values())
     output = ','.join(output_names)
@@ -88,8 +93,10 @@ def onnx2openvino(input_info: Dict[str, Union[List[int], torch.Size]],
               f'--output_dir="{work_dir}" ' \
               f'--output="{output}" ' \
               f'--input="{input_names}" ' \
-              f'--input_shape="{input_shapes}" ' \
-              f'--disable_fusing '
+              f'--input_shape="{input_shapes}" '
+    if mo_options is not None:
+        mo_args += mo_options.get_options()
+
     command = f'{mo_command} {mo_args}'
 
     logger = get_root_logger()

mmdeploy/backend/openvino/utils.py

@@ -0,0 +1,45 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List, Optional, Union
+
+
+class ModelOptimizerOptions:
+    """A class to make it easier to support additional arguments for the Model
+    Optimizer that can be passed through the deployment configuration.
+
+    Example:
+        >>> deploy_cfg = load_config(deploy_cfg_path)
+        >>> mo_options = deploy_cfg.get('mo_options', None)
+        >>> mo_options = ModelOptimizerOptions(mo_options)
+        >>> mo_args = mo_options.get_options()
+    """
+
+    def __init__(self,
+                 mo_options: Optional[Dict[str, Union[Dict, List]]] = None):
+        self.args = ''
+        self.flags = ''
+        if mo_options is not None:
+            self.args = self.__parse_args(mo_options)
+            self.flags = self.__parse_flags(mo_options)
+
+    def __parse_args(self, mo_options: Dict[str, Union[Dict, List]]) -> str:
+        """Parses a dictionary with arguments into a string."""
+        mo_args_str = ''
+        if 'args' in mo_options:
+            for key, value in mo_options['args'].items():
+                value_str = f'"{value}"' if isinstance(value, list) else value
+                mo_args_str += f'{key}={value_str} '
+        return mo_args_str
+
+    def __parse_flags(self, mo_options: Dict[str, Union[Dict, List]]) -> str:
+        """Parses a list with flags into a string."""
+        mo_flags_str = ''
+        if 'flags' in mo_options:
+            mo_flags_str += ' '.join(mo_options['flags'])
+        return mo_flags_str
+
+    def get_options(self) -> str:
+        """Returns a string with additional arguments for the Model Optimizer.
+
+        If there are no additional arguments, it will return an empty string.
+        """
+        return self.args + self.flags

mmdeploy/utils/test.py

@@ -487,15 +487,17 @@ def get_backend_outputs(ir_file_path: str,
         import mmdeploy.apis.openvino as openvino_apis
         if not openvino_apis.is_available():
             return None
+        from mmdeploy.apis.openvino import get_mo_options_from_cfg
         openvino_work_dir = tempfile.TemporaryDirectory().name
         openvino_file_path = openvino_apis.get_output_model_file(
             ir_file_path, openvino_work_dir)
         input_info = {
             name: value.shape
             for name, value in flatten_model_inputs.items()
         }
+        mo_options = get_mo_options_from_cfg(deploy_cfg)
         openvino_apis.onnx2openvino(input_info, output_names, ir_file_path,
-                                    openvino_work_dir)
+                                    openvino_work_dir, mo_options)
         backend_files = [openvino_file_path]
         backend_feats = flatten_model_inputs
         device = 'cpu'

tests/test_apis/test_onnx2openvino.py

@@ -60,9 +60,28 @@ def get_outputs(pytorch_model, openvino_model_path, input, input_name,
     return output_pytorch, openvino_output
 
 
+def get_base_deploy_cfg():
+    deploy_cfg = mmcv.Config(dict(backend_config=dict(type='openvino')))
+    return deploy_cfg
+
+
+def get_deploy_cfg_with_mo_args():
+    deploy_cfg = mmcv.Config(
+        dict(
+            backend_config=dict(
+                type='openvino',
+                mo_options=dict(
+                    args={'--data_type': 'FP32'}, flags=['--disable_fusing'
+                                                         ]))))
+    return deploy_cfg
+
+
+@pytest.mark.parametrize('get_deploy_cfg',
+                         [get_base_deploy_cfg, get_deploy_cfg_with_mo_args])
 @backend_checker(Backend.OPENVINO)
-def test_onnx2openvino():
-    from mmdeploy.apis.openvino import get_output_model_file, onnx2openvino
+def test_onnx2openvino(get_deploy_cfg):
+    from mmdeploy.apis.openvino import (get_mo_options_from_cfg,
+                                        get_output_model_file, onnx2openvino)
     pytorch_model = TestModel().eval()
     export_img = torch.rand([1, 3, 8, 8])
     onnx_file = tempfile.NamedTemporaryFile(suffix='.onnx').name
@@ -74,7 +93,10 @@ def test_onnx2openvino():
     input_info = {input_name: export_img.shape}
     output_names = [output_name]
     openvino_dir = tempfile.TemporaryDirectory().name
-    onnx2openvino(input_info, output_names, onnx_file, openvino_dir)
+    deploy_cfg = get_deploy_cfg()
+    mo_options = get_mo_options_from_cfg(deploy_cfg)
+    onnx2openvino(input_info, output_names, onnx_file, openvino_dir,
+                  mo_options)
     openvino_model_path = get_output_model_file(onnx_file, openvino_dir)
     assert osp.exists(openvino_model_path), \
         'The file (.xml) for OpenVINO IR has not been created.'

tools/deploy.py

@@ -222,17 +222,20 @@ def main():
             'OpenVINO is not available, please install OpenVINO first.'
 
         from mmdeploy.apis.openvino import (get_input_info_from_cfg,
+                                            get_mo_options_from_cfg,
                                             get_output_model_file,
                                             onnx2openvino)
         openvino_files = []
         for onnx_path in ir_files:
             model_xml_path = get_output_model_file(onnx_path, args.work_dir)
             input_info = get_input_info_from_cfg(deploy_cfg)
             output_names = get_ir_config(deploy_cfg).output_names
+            mo_options = get_mo_options_from_cfg(deploy_cfg)
             create_process(
                 f'onnx2openvino with {onnx_path}',
                 target=onnx2openvino,
-                args=(input_info, output_names, onnx_path, args.work_dir),
+                args=(input_info, output_names, onnx_path, args.work_dir,
+                      mo_options),
                 kwargs=dict(),
                 ret_value=ret_value)
             openvino_files.append(model_xml_path)