diff --git a/docs/en/07-developer-guide/support_new_backend.md b/docs/en/07-developer-guide/support_new_backend.md index 1513ca42c7..0357aee2d5 100644 --- a/docs/en/07-developer-guide/support_new_backend.md +++ b/docs/en/07-developer-guide/support_new_backend.md @@ -6,235 +6,259 @@ MMDeploy supports a number of backend engines. We welcome the contribution of ne Before contributing the codes, there are some requirements for the new backend that need to be checked: -- The backend must support ONNX as IR. +- The backend must support ONNX or Torchscript as IR. - If the backend requires model files or weight files other than a ".onnx" file, a conversion tool that converts the ".onnx" file to model files and weight files is required. The tool can be a Python API, a script, or an executable program. - It is highly recommended that the backend provides a Python interface to load the backend files and inference for validation. -## Support backend conversion +There are a lot of backends in `mmdeploy/backend`. Feel free to read the codes if you meet any problems. + +## Create BackendParam -The backends in MMDeploy must support the ONNX. The backend loads the ".onnx" file directly, or converts the ".onnx" to its own format using the conversion tool. In this section, we will introduce the steps to support backend conversion. +1. `BaseBackendParam` is a dataclass that we used to package everything parameters we need to do the conversion or inference. -1. Add backend constant in `mmdeploy/utils/constants.py` that denotes the name of the backend. + It is highly recommend to add google style docstring for your param. We will generate arguments for the `ArgumentParser` to ease the console tools. - **Example**: + ```python + # mmdeploy/backend/newengine/backend_manager.py + from ..base import BaseBackendParam + class NewEngineParam(BaseBackendParam): + """Your first backend parameters. + + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + """ + + work_dir:str = None + # FileNameDescriptor will add postfix to file name if it does not has one + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.model') + ``` + +2. Add a method to tell us the backend model names. + + ```python + # mmdeploy/backend/newengine/backend_manager.py + class NewEngineParam(BaseBackendParam): - ```Python - # mmdeploy/utils/constants.py + ... - class Backend(AdvancedEnum): - # Take TensorRT as an example - TENSORRT = 'tensorrt' + def get_model_files(self) -> Union[List, str]: + """get the model files.""" + return osp.join(self.work_dir, self.file_name) ``` -2. Add a corresponding package (a folder with `__init__.py`) in `mmdeploy/backend/`. For example, `mmdeploy/backend/tensorrt`. In the `__init__.py`, there must be a function named `is_available` which checks if users have installed the backend library. If the check is passed, then the remaining files of the package will be loaded. +## Check environment - **Example**: +A backend manager is the entry to everything about the backend. Assume your backend convert model from onnx. You can create the manager below: - ```Python - # mmdeploy/backend/tensorrt/__init__.py +```python +# mmdeploy/backend/newengine/backend_manager.py - def is_available(): - return importlib.util.find_spec('tensorrt') is not None +from mmdeploy.ir.onnx import ONNXParam +from ..base import BaseBackendManager +@BACKEND_MANAGERS.register('newengine', param=NewEngineParam, ir_param=ONNXParam) +class NewEngineManager(BaseBackendManager): +``` +If you decide to contribute the backend manager to MMDeploy, Do not forget to add enumrate in `mmdeploy/utils/constants.py` - if is_available(): - from .utils import from_onnx, load, save - from .wrapper import TRTWrapper +```Python +# mmdeploy/utils/constants.py - __all__ = [ - 'from_onnx', 'save', 'load', 'TRTWrapper' - ] - ``` +class Backend(AdvancedEnum): + # Take TensorRT as an example + NEWENGINE = 'newengine' +``` -3. Create a config file in `configs/_base_/backends` (e.g., `configs/_base_/backends/tensorrt.py`). If the backend just takes the '.onnx' file as input, the new config can be simple. The config of the backend only consists of one field denoting the name of the backend (which should be same as the name in `mmdeploy/utils/constants.py`). +Before we do anything with the backend. We want to make sure everything is fine. Let's add some method to check the environment. - **Example**: +- `is_available` return bool indicate that the backend manager is available on current device. +- `get_version` return the backend version information. +- `check_env` provide detail information about the backend. - ```python - backend_config = dict(type='onnxruntime') - ``` +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): - If the backend requires other files, then the arguments for the conversion from ".onnx" file to backend files should be included in the config file. + ... - **Example:** + @classmethod + def is_available(cls, with_custom_ops: bool = False) -> bool: + return my_backend_is_available() - ```Python + @classmethod + def get_version(cls) -> str: + return my_backend_version() - backend_config = dict( - type='tensorrt', - common_config=dict( - fp16_mode=False, max_workspace_size=0)) - ``` + @classmethod + def check_env(cls, log_callback: Callable = lambda _: _) -> str: + log_callback('Check env of your backend!') + return super().check_env(log_callback=log_callback) - After possessing a base backend config file, you can easily construct a complete deploy config through inheritance. Please refer to our [config tutorial](../02-how-to-run/write_config.md) for more details. Here is an example: +``` - ```Python - _base_ = ['../_base_/backends/onnxruntime.py'] +## Support backend conversion - codebase_config = dict(type='mmcls', task='Classification') - onnx_config = dict(input_shape=None) - ``` +Most backend has it's own serialize format. To support the conversion, Two method is required in backend manager: -4. If the backend requires model files or weight files other than a ".onnx" file, create a `onnx2backend.py` file in the corresponding folder (e.g., create `mmdeploy/backend/tensorrt/onnx2tensorrt.py`). Then add a conversion function `onnx2backend` in the file. The function should convert a given ".onnx" file to the required backend files in a given work directory. There are no requirements on other parameters of the function and the implementation details. You can use any tools for conversion. Here are some examples: +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): - **Use Python script:** + ... - ```Python - def onnx2openvino(input_info: Dict[str, Union[List[int], torch.Size]], - output_names: List[str], onnx_path: str, work_dir: str): + @classmethod + def to_backend(cls, ir_model: str, *args, **kwargs): + # convert your model here - input_names = ','.join(input_info.keys()) - input_shapes = ','.join(str(list(elem)) for elem in input_info.values()) - output = ','.join(output_names) + @classmethod + def to_backend_from_param(cls, ir_model: str, param: NewEngineParam): + # convert the model with the backend param you have just defined +``` - mo_args = f'--input_model="{onnx_path}" '\ - f'--output_dir="{work_dir}" ' \ - f'--output="{output}" ' \ - f'--input="{input_names}" ' \ - f'--input_shape="{input_shapes}" ' \ - f'--disable_fusing ' - command = f'mo.py {mo_args}' - mo_output = run(command, stdout=PIPE, stderr=PIPE, shell=True, check=True) - ``` +`to_backend` convert the model to the your backend. There is no limitation on the arguments of the method. It is up to you. - **Use executable program:** +`to_backend_from_param` accept a serialized IR file and the backend param you have just defined. You can extract fields from the backend param and convert model with `to_backend`. - ```Python - def onnx2ncnn(onnx_path: str, work_dir: str): - onnx2ncnn_path = get_onnx2ncnn_path() - save_param, save_bin = get_output_model_file(onnx_path, work_dir) - call([onnx2ncnn_path, onnx_path, save_param, save_bin])\ - ``` +## Support backend inference -5. Define APIs in a new package in `mmdeploy/apis`. +It would be cool if we can evaluate the backend model with python. Create a backend wrapper so we can perform inference with the backend and hide the detail. The inputs/outputs of the wrapper is a `dict` of `torch.Tensor`. - **Example:** +```python +# mmdeploy/backend/newengine/wrapper.py +from mmdeploy.utils import Backend +from ..base import BACKEND_WRAPPER, BaseWrapper - ```Python - # mmdeploy/apis/ncnn/__init__.py +@BACKEND_WRAPPER.register_module(Backend.NEWENGINE.value) +class NewEngineWrapper(BaseWrapper): - from mmdeploy.backend.ncnn import is_available + def __init__(self, backend_file, other_arguments): + # initialize the backend model here - __all__ = ['is_available'] + def forward(self, inputs) -> Dict[str, Tensor]: + # perform inference here +``` - if is_available(): - from mmdeploy.backend.ncnn.onnx2ncnn import (onnx2ncnn, - get_output_model_file) - __all__ += ['onnx2ncnn', 'get_output_model_file'] - ``` +Once you have a backend wrapper, add method in manager to build it. - Create a backend manager class which derive from `BaseBackendManager`, implement its `to_backend` static method. +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): - **Example:** + ... - ```Python @classmethod - def to_backend(cls, - ir_files: Sequence[str], - deploy_cfg: Any, - work_dir: str, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: - return ir_files - ``` + def build_wrapper(cls, backend_file, other_arguments): + from .wrapper import NewEngineWrapper + return NewEngineWrapper(backend_file, other_arguments) -6. Convert the models of OpenMMLab to backends (if necessary) and inference on backend engine. If you find some incompatible operators when testing, you can try to rewrite the original model for the backend following the [rewriter tutorial](support_new_model.md) or add custom operators. + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + backend_file = get_backend_file_from_param(param) + other_arguments = get_other_arguments_from_param(param) + return cls.build_wrapper(backend_file, other_arguments) +``` -7. Add docstring and unit tests for new code :). +## Console argument parser -## Support backend inference +What if you want to use the backend manager as a console tool. You can implement `parse_args` in the backend manager: -Although the backend engines are usually implemented in C/C++, it is convenient for testing and debugging if the backend provides Python inference interface. We encourage the contributors to support backend inference in the Python interface of MMDeploy. In this section we will introduce the steps to support backend inference. +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): -1. Add a file named `wrapper.py` to corresponding folder in `mmdeploy/backend/{backend}`. For example, `mmdeploy/backend/tensorrt/wrapper.py`. This module should implement and register a wrapper class that inherits the base class `BaseWrapper` in `mmdeploy/backend/base/base_wrapper.py`. + ... - **Example:** + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): - ```Python - from mmdeploy.utils import Backend - from ..base import BACKEND_WRAPPER, BaseWrapper + # setup parser + parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + NewEngineParam.add_arguments(parser) - @BACKEND_WRAPPER.register_module(Backend.TENSORRT.value) - class TRTWrapper(BaseWrapper): - ``` + parsed_args = parser.parse_args(args) -2. The wrapper class can initialize the engine in `__init__` function and inference in `forward` function. Note that the `__init__` function must take a parameter `output_names` and pass it to base class to determine the orders of output tensors. The input and output variables of `forward` should be dictionaries denoting the name and value of the tensors. + yield parsed_args -3. For the convenience of performance testing, the class should define a "execute" function that only calls the inference interface of the backend engine. The `forward` function should call the "execute" function after preprocessing the data. + # convert model + param = NewEngineParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name) - **Example:** + cls.to_backend_from_param(parsed_args.onnx_path, param) +``` - ```Python - from mmdeploy.utils import Backend - from mmdeploy.utils.timer import TimeCounter - from ..base import BACKEND_WRAPPER, BaseWrapper +`NewEngineParam.add_arguments(parser)` would add the arguments to the parser according to the docstring. You can create your console tool as: - @BACKEND_WRAPPER.register_module(Backend.ONNXRUNTIME.value) - class ORTWrapper(BaseWrapper): +```python +# console_tool.py +from ... import NewEngineManager - def __init__(self, - onnx_file: str, - device: str, - output_names: Optional[Sequence[str]] = None): - # Initialization - # ... - super().__init__(output_names) +if __name__ == '__main__' + NewEngineManager.main() +``` - def forward(self, inputs: Dict[str, - torch.Tensor]) -> Dict[str, torch.Tensor]: - # Fetch data - # ... +```bash +python console_tool.py -h + +# usage: test_dataclass.py convert [-h] --onnx-path ONNX_PATH +# [--work-dir WORK_DIR] [--file-name FILE_NAME] +# +# build TensorRTParam +# +# optional arguments: +# -h, --help show this help message and exit +# --onnx-path ONNX_PATH +# ONNX model path. +# --work-dir WORK_DIR The working directory. +# --file-name FILE_NAME +# File name of the serialized model. Postfix will +# beadded automatically. - self.__ort_execute(self.io_binding) +``` - # Postprocess data - # ... +## Unit Test - @TimeCounter.count_time('onnxruntime') - def __ort_execute(self, io_binding: ort.IOBinding): - # Only do the inference - self.sess.run_with_iobinding(io_binding) - ``` +Develop the unit test can help us test and maintain the backend support. Read scripts in `tests/test_backend` for model details. -4. Create a backend manager class which derive from `BaseBackendManager`, implement its `build_wrapper` static method. - - **Example:** - - ```Python - @BACKEND_MANAGERS.register('onnxruntime') - class ONNXRuntimeManager(BaseBackendManager): - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - from .wrapper import ORTWrapper - return ORTWrapper( - onnx_file=backend_files[0], - device=device, - output_names=output_names) - ``` +## Deploy config support -5. Add docstring and unit tests for new code :). +MMDeploy provide config files to describe the task you want to perform. Our main entry `tools/deploy.py` require the config file to perform the conversion and inference. If you hope the new backend can be use by `tools/deploy.py`, you would need a config too. -## Support new backends using MMDeploy as a third party +The config is a dictionary that composed of `model_config`, `ir_config`, `backend_config`... And they can be inherited by `__base__ = [...]`. +Create a config for the codebase you want to use the new backend. mmclassification as example: -Previous parts show how to add a new backend in MMDeploy, which requires changing its source codes. However, if we treat MMDeploy as a third party, the methods above are no longer efficient. To this end, adding a new backend requires us pre-install another package named `aenum`. We can install it directly through `pip install aenum`. +```python +_base_ = ['./classification_dynamic.py'] +codebase_config = dict(type='mmcls', task='Classification') +onnx_config = dict(input_shape=None) +backend_config = dict(type='newengine') +``` + +Read [config tutorial](../02-how-to-run/write_config.md) for more detail about the config file. -After installing `aenum` successfully, we can use it to add a new backend through: +Then add a convert tool in the backend manager: ```python -from mmdeploy.utils.constants import Backend -from aenum import extend_enum +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): + + ... + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> NewEngineParam: -try: - Backend.get('backend_name') -except Exception: - extend_enum(Backend, 'BACKEND', 'backend_name') + # create NewEngineParam with the config ``` -We can run the codes above before we use the rewrite logic of MMDeploy. +Now we have finish all steps. Enjoy the new backend and MMDeploy! diff --git a/docs/en/experimental/onnx_optimizer.md b/docs/en/experimental/onnx_optimizer.md index a40939d183..ab05ac27a7 100644 --- a/docs/en/experimental/onnx_optimizer.md +++ b/docs/en/experimental/onnx_optimizer.md @@ -21,9 +21,9 @@ cmake --build . -- -j$(nproc) && cmake --install . ```python # import model_to_graph_custom_optimizer so we can hijack onnx.export -from mmdeploy.apis.onnx.optimizer import model_to_graph__custom_optimizer # noqa +from mmdeploy.ir.onnx.optimizer import model_to_graph__custom_optimizer # noqa from mmdeploy.core import RewriterContext -from mmdeploy.apis.onnx.passes import optimize_onnx +from mmdeploy.ir.onnx.passes import optimize_onnx # load you model here model = create_model() diff --git a/docs/zh_cn/07-developer-guide/support_new_backend.md b/docs/zh_cn/07-developer-guide/support_new_backend.md index 91b1cf3393..fad8c12729 100644 --- a/docs/zh_cn/07-developer-guide/support_new_backend.md +++ b/docs/zh_cn/07-developer-guide/support_new_backend.md @@ -10,232 +10,255 @@ MMDeploy 支持了许多后端推理引擎，但我们依然非常欢迎新后 - 如果后端需要“.onnx”文件以外的模型文件或权重文件，则需要添加将“.onnx”文件转换为模型文件或权重文件的转换工具，该工具可以是 Python API、脚本或可执行程序。 - 强烈建议新后端可提供 Python 接口来加载后端文件和推理以进行验证。 -## 支持后端转换 +`mmdeploy/backend` 目录下有许多已经接入的后端。如果在实现时存在任何困难，可以参考其中的代码实现。 + +## 创建 BackendParam + +1. `BaseBackendParam` 是一个 dataclass 类，我们将模型转换与推理需要的所有数据打包在其中，方便后续的操作 -MMDeploy 中的后端必须支持 ONNX，因此后端能直接加载“.onnx”文件，或者使用转换工具将“.onnx”转换成自己的格式。在本节中，我们将介绍支持后端转换的步骤。 + 我们非常推荐您给 BackendParam 的实现提供一个 google 风格的 docsting，我们会根据 docstring 的内容生成一个命令行参数解析器，帮助创建命令行工具 -1. 在 `mmdeploy/utils/constants.py` 文件中添加新推理后端变量，以表示支持的后端名称。 + ```python + # mmdeploy/backend/newengine/backend_manager.py + from ..base import BaseBackendParam + class NewEngineParam(BaseBackendParam): + """Your first backend parameters. + + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + """ + + work_dir:str = None + # FileNameDescriptor will add postfix to file name if it does not has one + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.model') + ``` - **示例**： +2. 给创建的 param 对象提供一个方法，查询转换后的文件名。如果转换后或生成多个文件，则返回一个 string 列表 - ```Python - # mmdeploy/utils/constants.py + ```python + # mmdeploy/backend/newengine/backend_manager.py + class NewEngineParam(BaseBackendParam): - class Backend(AdvancedEnum): - # 以现有的TensorRT为例 - TENSORRT = 'tensorrt' + ... + + def get_model_files(self) -> Union[List, str]: + """get the model files.""" + return osp.join(self.work_dir, self.file_name) ``` -2. 在 `mmdeploy/backend/` 目录下添加相应的库(一个包括 `__init__.py` 的文件夹),例如， `mmdeploy/backend/tensorrt` 。在 `__init__.py` 中，必须有一个名为 `is_available` 的函数检查用户是否安装了后端库。如果检查通过，则将加载库的剩余文件。 +## 环境检查 - **例子**: +每个接入的后端以 BackendManager 对象作为入口，它提供了包括转换、推理、环境检查等功能。 - ```Python - # mmdeploy/backend/tensorrt/__init__.py +```python +# mmdeploy/backend/newengine/backend_manager.py - def is_available(): - return importlib.util.find_spec('tensorrt') is not None +from mmdeploy.ir.onnx import ONNXParam +from ..base import BaseBackendManager +@BACKEND_MANAGERS.register('newengine', param=NewEngineParam, ir_param=ONNXParam) +class NewEngineManager(BaseBackendManager): +``` +如果你希望将接入的后端贡献给 MMDeploy，那么可以在 `mmdeploy/utils/constants.py` 中添加如下代码并接受我们真挚的感谢 - if is_available(): - from .utils import from_onnx, load, save - from .wrapper import TRTWrapper +```Python +# mmdeploy/utils/constants.py - __all__ = [ - 'from_onnx', 'save', 'load', 'TRTWrapper' - ] - ``` +class Backend(AdvancedEnum): + # Take TensorRT as an example + NEWENGINE = 'newengine' +``` -3. 在 `configs/_base_/backends` 目录中创建一个配置文件(例如， `configs/_base_/backends/tensorrt.py` )。如果新后端引擎只是将“.onnx”文件作为输入，那么新的配置可以很简单,对应配置只需包含一个表示后端名称的字段(但也应该与 `mmdeploy/utils/constants.py` 中的名称相同)。 +在我们真正开始进行功能的接入之前，首先应该对当前环境进行检查以确保后续运算可以正确进行。因此我们需要如下接口： - **例子** +- `is_available` 返回 bool 值，表示该后端在当前环境下可用 +- `get_version` 返回该后端的版本信息 +- `check_env` 提供更详细的当前环境信息，用于帮助用户配置环境 - ```python - backend_config = dict(type='tensorrt') - ``` +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): - 但如果后端需要其他文件，则从“.onnx”文件转换为后端文件所需的参数也应包含在配置文件中。 + ... - **例子** + @classmethod + def is_available(cls, with_custom_ops: bool = False) -> bool: + return my_backend_is_available() - ```Python + @classmethod + def get_version(cls) -> str: + return my_backend_version() - backend_config = dict( - type='tensorrt', - common_config=dict( - fp16_mode=False, max_workspace_size=0)) - ``` + @classmethod + def check_env(cls, log_callback: Callable = lambda _: _) -> str: + log_callback('Check env of your backend!') + return super().check_env(log_callback=log_callback) - 在拥有一个基本的后端配置文件后，您已经可以通过继承轻松构建一个完整的部署配置。有关详细信息，请参阅我们的[配置教程](../02-how-to-run/write_config.md)。下面是一个例子： +``` - ```Python - _base_ = ['../_base_/backends/tensorrt.py'] +## 支持后端转换 - codebase_config = dict(type='mmcls', task='Classification') - onnx_config = dict(input_shape=None) - ``` +多数后端都会有自己的序列化模型格式，为了支持从中间表示到该格式的转换，需要提供两个函数： -4. 如果新后端需要模型文件或权重文件而不是“.onnx”文件，则需要在相应的文件夹中创建一个 `onnx2backend.py` 文件(例如,创建 `mmdeploy/backend/tensorrt/onnx2tensorrt.py` )。然后在文件中添加一个转换函数`onnx2backend`。该函数应将给定的“.onnx”文件转换为给定工作目录中所需的后端文件。对函数的其他参数和实现细节没有要求，您可以使用任何工具进行转换。下面有些例子： +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): - **使用python脚本** + ... - ```Python - def onnx2openvino(input_info: Dict[str, Union[List[int], torch.Size]], - output_names: List[str], onnx_path: str, work_dir: str): + @classmethod + def to_backend(cls, ir_model: str, *args, **kwargs): + # convert your model here - input_names = ','.join(input_info.keys()) - input_shapes = ','.join(str(list(elem)) for elem in input_info.values()) - output = ','.join(output_names) + @classmethod + def to_backend_from_param(cls, ir_model: str, param: NewEngineParam): + # convert the model with the backend param you have just defined +``` - mo_args = f'--input_model="{onnx_path}" '\ - f'--output_dir="{work_dir}" ' \ - f'--output="{output}" ' \ - f'--input="{input_names}" ' \ - f'--input_shape="{input_shapes}" ' \ - f'--disable_fusing ' - command = f'mo.py {mo_args}' - mo_output = run(command, stdout=PIPE, stderr=PIPE, shell=True, check=True) - ``` +`to_backend` 用来将 IR 模型转换成后端需要的格式，我们不对输入参数做太多限制，可以根据后端的需要自由配置。 +`to_backend_from_param` 使用上面章节实现的 BackendParam 对象来实现转换。可以从 param 中提取数据然后调用 `to_backend` 以复用代码。 - **使用可执行文件** +## 支持后端推理 - ```Python - def onnx2ncnn(onnx_path: str, work_dir: str): - onnx2ncnn_path = get_onnx2ncnn_path() - save_param, save_bin = get_output_model_file(onnx_path, work_dir) - call([onnx2ncnn_path, onnx_path, save_param, save_bin])\ - ``` +如果希望可以使用 python 进行精度验证，那么就需要实现一个 Wrapper 对象和对应的构建函数。Wrapper 对象对后端推理细节进行了封装，用户可以像使用 PyTorch 模型那样使用后端接口。Wrapper 的输入输出为 Tensor 的 dict 对象。 -5. 在 `mmdeploy/apis` 中创建新后端库并声明对应 APIs +```python +# mmdeploy/backend/newengine/wrapper.py +from mmdeploy.utils import Backend +from ..base import BACKEND_WRAPPER, BaseWrapper - **例子** +@BACKEND_WRAPPER.register_module(Backend.NEWENGINE.value) +class NewEngineWrapper(BaseWrapper): - ```Python - # mmdeploy/apis/ncnn/__init__.py + def __init__(self, backend_file, other_arguments): + # initialize the backend model here - from mmdeploy.backend.ncnn import is_available + def forward(self, inputs) -> Dict[str, Tensor]: + # perform inference here +``` - __all__ = ['is_available'] +实现了 Wrapper 以后，就可以在 backend manager 中添加构建函数： - if is_available(): - from mmdeploy.backend.ncnn.onnx2ncnn import (onnx2ncnn, - get_output_model_file) - __all__ += ['onnx2ncnn', 'get_output_model_file'] - ``` +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): - 从 BaseBackendManager 派生类，实现 `to_backend` 类方法。 + ... - **例子** + @classmethod + def build_wrapper(cls, backend_file, other_arguments): + from .wrapper import NewEngineWrapper + return NewEngineWrapper(backend_file, other_arguments) - ```Python - @classmethod - def to_backend(cls, - ir_files: Sequence[str], - deploy_cfg: Any, - work_dir: str, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: - return ir_files - ``` + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + backend_file = get_backend_file_from_param(param) + other_arguments = get_other_arguments_from_param(param) + return cls.build_wrapper(backend_file, other_arguments) +``` -6. 将 OpenMMLab 的模型转换后(如有必要)并在后端引擎上进行推理。如果在测试时发现一些不兼容的算子，可以尝试按照[重写器教程](support_new_model.md)为后端重写原始模型或添加自定义算子。 +## 命令行工具 -7. 为新后端引擎代码添加相关注释和单元测试:). +如果希望将接入的后端作为一个命令行工具使用，可以实现 `parse_args` 接口： -## 支持后端推理 +```python +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): -尽管后端引擎通常用C/C++实现，但如果后端提供Python推理接口，则测试和调试非常方便。我们鼓励贡献者在MMDeploy的Python接口中支持新后端推理。在本节中，我们将介绍支持后端推理的步骤。 + ... -1. 添加一个名为 `wrapper.py` 的文件到 `mmdeploy/backend/{backend}` 中相应后端文件夹。例如， `mmdeploy/backend/tensorrt/wrapper` 。此模块应实现并注册一个封装类，该类继承 `mmdeploy/backend/base/base_wrapper.py` 中的基类 `BaseWrapper` 。 + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): - **例子** + # setup parser + parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + NewEngineParam.add_arguments(parser) - ```Python - from mmdeploy.utils import Backend - from ..base import BACKEND_WRAPPER, BaseWrapper + parsed_args = parser.parse_args(args) - @BACKEND_WRAPPER.register_module(Backend.TENSORRT.value) - class TRTWrapper(BaseWrapper): - ``` + yield parsed_args -2. 封装类可以在函数 `__init__` 中初始化引擎以及在 `forward` 函数中进行推理。请注意，该 `__init__` 函数必须接受一个参数 `output_names` 并将其传递给基类以确定输出张量的顺序。其中 `forward` 输入和输出变量应表示tensors的名称和值的字典。 + # convert model + param = NewEngineParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name) -3. 为了方便性能测试，该类应该定义一个 `execute` 函数，只调用后端引擎的推理接口。该 `forward` 函数应在预处理数据后调用 `execute` 函数。 + cls.to_backend_from_param(parsed_args.onnx_path, param) +``` - **例子** +`NewEngineParam.add_arguments(parser)` 会根据之前在 `NewEngineParam` 中添加的 docstring 信息自动生成解析器的 arguments。方便我们更快实现功能。 - ```Python - from mmdeploy.utils import Backend - from mmdeploy.utils.timer import TimeCounter - from ..base import BACKEND_WRAPPER, BaseWrapper +我们只要实现下面几行代码，就可以完成该工具： - @BACKEND_WRAPPER.register_module(Backend.ONNXRUNTIME.value) - class ORTWrapper(BaseWrapper): +```python +# console_tool.py +from ... import NewEngineManager - def __init__(self, - onnx_file: str, - device: str, - output_names: Optional[Sequence[str]] = None): - # Initialization - # - # ... - super().__init__(output_names) +if __name__ == '__main__' + NewEngineManager.main() +``` - def forward(self, inputs: Dict[str, - torch.Tensor]) -> Dict[str, torch.Tensor]: - # Fetch data - # ... +使用效果如下 + +```bash +python console_tool.py -h + +# usage: test_dataclass.py convert [-h] --onnx-path ONNX_PATH +# [--work-dir WORK_DIR] [--file-name FILE_NAME] +# +# build TensorRTParam +# +# optional arguments: +# -h, --help show this help message and exit +# --onnx-path ONNX_PATH +# ONNX model path. +# --work-dir WORK_DIR The working directory. +# --file-name FILE_NAME +# File name of the serialized model. Postfix will +# beadded automatically. - self.__ort_execute(self.io_binding) +``` - # Postprocess data - # ... +## 单元测试 - @TimeCounter.count_time('onnxruntime') - def __ort_execute(self, io_binding: ort.IOBinding): - # Only do the inference - self.sess.run_with_iobinding(io_binding) - ``` +开发单元测试是一个好习惯，可以为功能维护以及更新带来便利。可以参考 `tests/test_backend` 添加自己的后端单元测试。 -4. 从 `BaseBackendManager` 派生接口类，实现 `build_wrapper` 静态方法 - - **例子** - - ```Python - @BACKEND_MANAGERS.register('onnxruntime') - class ONNXRuntimeManager(BaseBackendManager): - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - from .wrapper import ORTWrapper - return ORTWrapper( - onnx_file=backend_files[0], - device=device, - output_names=output_names) - ``` +## deploy.py 支持 -5. 为新后端引擎代码添加相关注释和单元测试 :). +当前 MMDeploy 的总接口为 `tools/deploy.py`。它需要一个配置文件来实现转换、推理任务。如果希望接口的后端能够使用该接口，那么后端应该提供自己的配置文件。 -## 将MMDeploy作为第三方库时添加新后端 +配置文件的写法可以参考 [config tutorial](../02-how-to-run/write_config.md)， 这里不做赘述。假设我们为新添加的后端提供了如下配置文件： -前面的部分展示了如何在 MMDeploy 中添加新的后端，这需要更改其源代码。但是，如果我们将 MMDeploy 视为第三方，则上述方法不再有效。为此，添加一个新的后端需要我们预先安装另一个名为 `aenum` 的包。我们可以直接通过`pip install aenum`进行安装。 +```python +_base_ = ['./classification_dynamic.py'] +codebase_config = dict(type='mmcls', task='Classification') +onnx_config = dict(input_shape=None) +backend_config = dict(type='newengine') +``` -成功安装 `aenum` 后，我们可以通过以下方式使用它来添加新的后端： +需要在 backend manager 对象中添加解析函数，通过 config 生成 BackendParam 对象 ```python -from mmdeploy.utils.constants import Backend -from aenum import extend_enum +# mmdeploy/backend/newengine/backend_manager.py +class NewEngineManager(BaseBackendManager): + + ... + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> NewEngineParam: -try: - Backend.get('backend_name') -except Exception: - extend_enum(Backend, 'BACKEND', 'backend_name') + # create NewEngineParam with the config ``` -我们可以在使用 MMDeploy 的重写逻辑之前运行上面的代码，这就完成了新后端的添加。 +如果能够正确配置上述步骤，那么恭喜你，你已经完成了后端接入，可以开始在 MMDeploy 中享受新的推理后端！ diff --git a/docs/zh_cn/experimental/onnx_optimizer.md b/docs/zh_cn/experimental/onnx_optimizer.md index a40939d183..ab05ac27a7 100644 --- a/docs/zh_cn/experimental/onnx_optimizer.md +++ b/docs/zh_cn/experimental/onnx_optimizer.md @@ -21,9 +21,9 @@ cmake --build . -- -j$(nproc) && cmake --install . ```python # import model_to_graph_custom_optimizer so we can hijack onnx.export -from mmdeploy.apis.onnx.optimizer import model_to_graph__custom_optimizer # noqa +from mmdeploy.ir.onnx.optimizer import model_to_graph__custom_optimizer # noqa from mmdeploy.core import RewriterContext -from mmdeploy.apis.onnx.passes import optimize_onnx +from mmdeploy.ir.onnx.passes import optimize_onnx # load you model here model = create_model() diff --git a/docs/zh_cn/tutorial/07_write_a_plugin.md b/docs/zh_cn/tutorial/07_write_a_plugin.md index f2bc5b5584..74e9e7f00c 100644 --- a/docs/zh_cn/tutorial/07_write_a_plugin.md +++ b/docs/zh_cn/tutorial/07_write_a_plugin.md @@ -490,7 +490,7 @@ engine = from_onnx( opt_shape=[1, 1, 512, 512], max_shape=[1, 1, 1024, 1024]))) -from mmdeploy.backend.tensorrt import TRTWrapper +from mmdeploy.backend.tensorrt.wrapper import TRTWrapper trt_model = TRTWrapper('srcnn3.engine', ['output']) diff --git a/mmdeploy/__main__.py b/mmdeploy/__main__.py new file mode 100644 index 0000000000..665ea97c26 --- /dev/null +++ b/mmdeploy/__main__.py @@ -0,0 +1,46 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import sys as _sys + +if __name__ == '__main__': + # args default to the system args + console_args = _sys.argv[1:] + + # extract help + help = False + if '-h' in console_args: + help = True + console_args.remove('-h') + if '--help' in console_args: + help = True + console_args.remove('--help') + + # add root parser + parser = argparse.ArgumentParser( + 'mmdeploy', description='MMDeploy Toolkit') + command_parsers = parser.add_subparsers(title='Commands', dest='command') + list_parser = command_parsers.add_parser( + 'list', help='List available backend and task.') + show_parser = command_parsers.add_parser( + 'show', help='Should information about the object.') + run_parser = command_parsers.add_parser( + 'run', help='Run console tools of backend or task.') + args, remain_args = parser.parse_known_args(console_args) + + # parse command + command = getattr(args, 'command', None) + + if help: + remain_args = ['--help'] + remain_args + if command == 'list': + from mmdeploy.tools.console import list_command + list_command(list_parser, remain_args) + elif command == 'show': + from mmdeploy.tools.console import show_command + show_command(list_parser, remain_args) + elif command == 'run': + from mmdeploy.tools.console import run_command + run_command(list_parser, remain_args) + else: + parser.print_help() + parser.exit() diff --git a/mmdeploy/apis/onnx/export.py b/mmdeploy/apis/onnx/export.py index 92a9002d8d..982af060a4 100644 --- a/mmdeploy/apis/onnx/export.py +++ b/mmdeploy/apis/onnx/export.py @@ -1,15 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -from copy import deepcopy -from functools import partial -from typing import Any, Dict, Optional, Sequence, Tuple, Union +from typing import Dict, Optional, Sequence, Tuple, Union import torch from mmdeploy.apis.core import PIPELINE_MANAGER -from mmdeploy.core import RewriterContext, patch_model -from mmdeploy.utils import IR, Backend, get_ir_config, get_root_logger -from .optimizer import * # noqa -from .passes import optimize_onnx +from mmdeploy.ir.onnx import ONNXManager +from mmdeploy.utils import Backend @PIPELINE_MANAGER.register_pipeline() @@ -70,75 +66,17 @@ def export(model: torch.nn.Module, """ output_path = output_path_prefix + '.onnx' - logger = get_root_logger() - logger.info(f'Export PyTorch model to ONNX: {output_path}.') - - def _add_or_update(cfg: dict, key: str, val: Any): - if key in cfg and isinstance(cfg[key], dict) and isinstance(val, dict): - cfg[key].update(val) - else: - cfg[key] = val - - context_info = deepcopy(context_info) deploy_cfg = context_info.pop('deploy_cfg', dict()) - ir_config = dict( - type='onnx', + ONNXManager.export( + model, + args, + output_path, input_names=input_names, output_names=output_names, opset_version=opset_version, dynamic_axes=dynamic_axes, verbose=verbose, - keep_initializers_as_inputs=keep_initializers_as_inputs) - _add_or_update(deploy_cfg, 'ir_config', ir_config) - ir = IR.get(get_ir_config(deploy_cfg)['type']) - if isinstance(backend, Backend): - backend = backend.value - backend_config = dict(type=backend) - _add_or_update(deploy_cfg, 'backend_config', backend_config) - - context_info['cfg'] = deploy_cfg - context_info['ir'] = ir - if 'backend' not in context_info: - context_info['backend'] = backend - if 'opset' not in context_info: - context_info['opset'] = opset_version - - # patch model - patched_model = patch_model(model, cfg=deploy_cfg, backend=backend, ir=ir) - - if 'onnx_custom_passes' not in context_info: - onnx_custom_passes = optimize_onnx if optimize else None - context_info['onnx_custom_passes'] = onnx_custom_passes - with RewriterContext(**context_info), torch.no_grad(): - # patch input_metas - if input_metas is not None: - assert isinstance( - input_metas, dict - ), f'Expect input_metas type is dict, get {type(input_metas)}.' - model_forward = patched_model.forward - - def wrap_forward(forward): - - def wrapper(*arg, **kwargs): - return forward(*arg, **kwargs) - - return wrapper - - patched_model.forward = wrap_forward(patched_model.forward) - patched_model.forward = partial(patched_model.forward, - **input_metas) - - torch.onnx.export( - patched_model, - args, - output_path, - export_params=True, - input_names=input_names, - output_names=output_names, - opset_version=opset_version, - dynamic_axes=dynamic_axes, - keep_initializers_as_inputs=keep_initializers_as_inputs, - verbose=verbose) - - if input_metas is not None: - patched_model.forward = model_forward + backend=backend, + const_args=input_metas, + rewrite_context=deploy_cfg, + optimize=optimize) diff --git a/mmdeploy/apis/openvino/utils.py b/mmdeploy/apis/openvino/utils.py index 41b2487213..af7c2ce239 100644 --- a/mmdeploy/apis/openvino/utils.py +++ b/mmdeploy/apis/openvino/utils.py @@ -3,7 +3,7 @@ import mmengine -from mmdeploy.backend.openvino import ModelOptimizerOptions +from mmdeploy.backend.openvino.utils import ModelOptimizerOptions from mmdeploy.utils import get_model_inputs from mmdeploy.utils.config_utils import get_backend_config, get_ir_config diff --git a/mmdeploy/apis/snpe/__init__.py b/mmdeploy/apis/snpe/__init__.py index 6f8febaec3..043e217095 100644 --- a/mmdeploy/apis/snpe/__init__.py +++ b/mmdeploy/apis/snpe/__init__.py @@ -1,6 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. -from mmdeploy.backend.snpe import from_onnx as _from_onnx from mmdeploy.backend.snpe import is_available +from mmdeploy.backend.snpe.onnx2dlc import from_onnx as _from_onnx from ..core import PIPELINE_MANAGER from_onnx = PIPELINE_MANAGER.register_pipeline()(_from_onnx) diff --git a/mmdeploy/apis/tensorrt/__init__.py b/mmdeploy/apis/tensorrt/__init__.py index d3010b37a7..bd53c09a97 100644 --- a/mmdeploy/apis/tensorrt/__init__.py +++ b/mmdeploy/apis/tensorrt/__init__.py @@ -5,8 +5,8 @@ __all__ = ['is_available'] if is_available(): - from mmdeploy.backend.tensorrt import from_onnx as _from_onnx - from mmdeploy.backend.tensorrt import load, save + from mmdeploy.backend.tensorrt.utils import from_onnx as _from_onnx + from mmdeploy.backend.tensorrt.utils import load, save from_onnx = PIPELINE_MANAGER.register_pipeline()(_from_onnx) __all__ += ['from_onnx', 'save', 'load'] try: diff --git a/mmdeploy/apis/torch_jit/trace.py b/mmdeploy/apis/torch_jit/trace.py index 0e1e33f5b3..1d9d28cf46 100644 --- a/mmdeploy/apis/torch_jit/trace.py +++ b/mmdeploy/apis/torch_jit/trace.py @@ -1,12 +1,10 @@ # Copyright (c) OpenMMLab. All rights reserved. -from copy import deepcopy -from functools import partial -from typing import Any, Dict, Optional, Sequence, Tuple, Union +from typing import Dict, Optional, Tuple, Union import torch -from mmdeploy.core import RewriterContext, patch_model -from mmdeploy.utils import IR, Backend, get_ir_config, get_root_logger +from mmdeploy.ir.torchscript import export +from mmdeploy.utils import Backend from ..core import PIPELINE_MANAGER @@ -27,9 +25,10 @@ def trace(func: torch.nn.Module, >>> func = create_model() >>> inputs = get_input_tensor() >>> - >>> jit_model = trace( + >>> trace( >>> func, >>> inputs, + >>> output_prefix, >>> backend='torchscript', >>> check_trace=False) >>> @@ -55,69 +54,23 @@ def trace(func: torch.nn.Module, Returns: torch.jit.TracedModule: The traced torch jit model. """ - logger = get_root_logger() - logger.info('Export PyTorch model to torchscript.') + if output_path_prefix is None: + from tempfile import NamedTemporaryFile + output_path = NamedTemporaryFile(suffix='.pth').name + else: + output_path = output_path_prefix + '.pth' - def _add_or_update(cfg: dict, key: str, val: Any): - if key in cfg and isinstance(cfg[key], dict) and isinstance(val, dict): - cfg[key].update(val) - else: - cfg[key] = val - - context_info = deepcopy(context_info) deploy_cfg = context_info.pop('deploy_cfg', dict()) - ir_config = dict(type='torchscript') - _add_or_update(deploy_cfg, 'ir_config', ir_config) - - if isinstance(backend, Backend): - backend = backend.value - backend_config = dict(type=backend) - _add_or_update(deploy_cfg, 'backend_config', backend_config) - - context_info['cfg'] = deploy_cfg - if 'backend' not in context_info: - context_info['backend'] = backend - elif context_info['backend'] != backend: - logger.warning( - f'Find backend {context_info["backend"]} in context_info.' - f' Expect {backend}.') - if 'ir' not in context_info: - context_info['ir'] = IR.TORCHSCRIPT - elif context_info['ir'] != backend: - logger.warning(f'Find ir {context_info["ir"]} in context_info.' - f' Expect {IR.TORCHSCRIPT}.') - - # patch model - if isinstance(func, torch.nn.Module): - ir = IR.get(get_ir_config(deploy_cfg)['type']) - func = patch_model(func, cfg=deploy_cfg, backend=backend, ir=ir) - - with RewriterContext(**context_info), torch.no_grad(): - - # patch input_metas - if input_metas is not None: - assert isinstance( - input_metas, dict - ), f'Expect input_metas type is dict, get {type(input_metas)}.' - model_forward = func.forward - func.forward = partial(func.forward, **input_metas) - - # for exporting models with weight that depends on inputs - func(*inputs) if isinstance(inputs, Sequence) \ - else func(inputs) - ts_model = torch.jit.trace( - func, - inputs, - check_trace=check_trace, - check_tolerance=check_tolerance) - - if input_metas is not None: - func.forward = model_forward - - # save model - if output_path_prefix is not None: - output_path = output_path_prefix + '.pt' - logger.info(f'Save PyTorch model: {output_path}.') - torch.jit.save(ts_model, output_path) + export( + func, + inputs, + output_path, + backend=backend, + rewrite_context=deploy_cfg, + check_trace=check_trace, + check_tolerance=check_tolerance, + const_args=input_metas) + + ts_model = torch.jit.load(output_path) return ts_model diff --git a/mmdeploy/apis/tvm/__init__.py b/mmdeploy/apis/tvm/__init__.py index fd4040592b..6c0b314f7d 100644 --- a/mmdeploy/apis/tvm/__init__.py +++ b/mmdeploy/apis/tvm/__init__.py @@ -5,8 +5,7 @@ __all__ = ['is_available', 'get_library_ext'] if is_available(): - from mmdeploy.backend.tvm import HDF5Dataset - from mmdeploy.backend.tvm import from_onnx as _from_onnx + from mmdeploy.backend.tvm.onnx2tvm import from_onnx as _from_onnx from_onnx = PIPELINE_MANAGER.register_pipeline()(_from_onnx) - __all__ += ['from_onnx', 'HDF5Dataset'] + __all__ += ['from_onnx'] diff --git a/mmdeploy/apis/utils/utils.py b/mmdeploy/apis/utils/utils.py index d7630e6637..00f59ee1b1 100644 --- a/mmdeploy/apis/utils/utils.py +++ b/mmdeploy/apis/utils/utils.py @@ -93,12 +93,38 @@ def to_backend(backend_name: str, Returns: Sequence[str]: Backend files. """ + import os.path as osp + from copy import deepcopy + from mmdeploy.backend.base import get_backend_manager + from mmdeploy.utils import get_model_inputs backend_mgr = get_backend_manager(backend_name) - return backend_mgr.to_backend( - ir_files=ir_files, - work_dir=work_dir, - deploy_cfg=deploy_cfg, - log_level=log_level, - device=device, - **kwargs) + + model_inputs = get_model_inputs(deploy_cfg) + assert model_inputs is None or len(model_inputs) == 0 or len( + model_inputs) == len(ir_files) + backend_files = [] + for idx, ir_file in enumerate(ir_files): + if isinstance(model_inputs, (list, tuple)) and len(model_inputs) > 0: + curr_deploy_cfg = deepcopy(deploy_cfg) + curr_deploy_cfg['backend_config']['model_inputs'] = [ + model_inputs[idx] + ] + else: + curr_deploy_cfg = deploy_cfg + + file_name = osp.splitext(osp.split(ir_file)[1])[0] + param = backend_mgr.build_param_from_config( + curr_deploy_cfg, + work_dir=work_dir, + backend_files=[file_name], + device=device, + **kwargs) + + backend_mgr.to_backend_from_param(ir_file, param) + backend_file = param.get_model_files() + if isinstance(backend_file, str): + backend_file = [backend_file] + backend_files += backend_file + + return backend_files diff --git a/mmdeploy/apis/vacc/__init__.py b/mmdeploy/apis/vacc/__init__.py index 5ab2116d98..b1ca23d719 100644 --- a/mmdeploy/apis/vacc/__init__.py +++ b/mmdeploy/apis/vacc/__init__.py @@ -6,7 +6,7 @@ if is_available(): try: - from mmdeploy.backend.vacc import from_onnx as _from_onnx + from mmdeploy.backend.vacc.onnx2vacc import from_onnx as _from_onnx from_onnx = PIPELINE_MANAGER.register_pipeline()(_from_onnx) __all__ += ['from_onnx'] except Exception: diff --git a/mmdeploy/backend/ascend/__init__.py b/mmdeploy/backend/ascend/__init__.py index 5b70bf8d8e..aeec8db578 100644 --- a/mmdeploy/backend/ascend/__init__.py +++ b/mmdeploy/backend/ascend/__init__.py @@ -1,14 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import AscendManager +from .backend_manager import AscendManager, AscendParam +from .onnx2ascend import AtcParam from .utils import update_sdk_pipeline _BackendManager = AscendManager - is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['update_sdk_pipeline', 'AscendManager'] - -if is_available(): - from .wrapper import AscendWrapper, Error - __all__ += ['AscendWrapper', 'Error'] +__all__ = ['update_sdk_pipeline', 'AtcParam', 'AscendParam', 'AscendManager'] diff --git a/mmdeploy/backend/ascend/backend_manager.py b/mmdeploy/backend/ascend/backend_manager.py index 165e3ba583..750814068a 100644 --- a/mmdeploy/backend/ascend/backend_manager.py +++ b/mmdeploy/backend/ascend/backend_manager.py @@ -1,36 +1,102 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os.path as osp -from typing import Any, Optional, Sequence +import re +from argparse import Action, ArgumentParser +from dataclasses import dataclass +from typing import Any, List, Optional, Sequence -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor, import_custom_modules) +from .onnx2ascend import AtcParam -@BACKEND_MANAGERS.register('ascend') -class AscendManager(BaseBackendManager): +class DynamicDimsAction(Action): + """dynamic dims argparse action.""" + + def __call__(self, parser, namespace, values, option_string=None): + """call action.""" + args = values + if isinstance(args, str): + args = [args] + + pattern = r'^[^\S\n]*((?P\w+):)?(?P([+|-]?\d+)(,[+|-]?\d+)*)$' # noqa + ret = dict() + for arg in args: + arg = [i.strip() for i in arg.split(';')] + for single_arg in arg: + if len(single_arg) == 0: + continue + m = re.match(pattern, single_arg) + if m is None: + raise ValueError(f'Can not parse value: {single_arg}') + input_name = m.group('input_name') + val = m.group('val') + val = val.split(',') + val = tuple(int(v) for v in val) + if input_name in ret: + raise NameError(f'value of `{input_name}` ' + 'has been assigned more than once.') + ret[input_name] = val + + setattr(namespace, self.dest, ret) + + +@dataclass +class AscendParam(BaseBackendParam): + """Ascend backend parameters. + + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + input_shapes (ShapeType): The Default shape of the inputs. + dynamic_batch_size (List[str]): Set dynamic batch size. + E.g.: "batchsize1 batchsize2 batchsize3" + dynamic_image_size (Dict[List[int]]): Set dynamic image size. + Separate multiple nodes with semicolons (;). + Use double quotation marks (") to enclose each argument. + E.g.: "input0:height0,width0;input1:height1,width1" + dynamic_dims (Dict[List[int]]): Set dynamic dims. + Separate multiple nodes with semicolons (;). + Use double quotation marks (") to enclose each argument. + E.g.: "input0:dims1_n1,dims1_n2;input1:dims2_n1,dims2_n2" + device (str): Inference device. + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.om') + dynamic_batch_size: List[str] = None + dynamic_image_size: List[List[int]] = None + dynamic_dims: List[List[int]] = None + + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str), ('Expect string work_dir, ' + f'got {self.work_dir}') + assert isinstance(self.file_name, str), ('Expect string file_name, ' + f'got {self.file_name}') + file_name = self.file_name + if not file_name.endswith('.om'): + file_name = file_name + '.om' + return osp.join(self.work_dir, file_name) @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + def add_argument(cls, parser: ArgumentParser, name: str, dtype: Any, + default: Any, desc: str): + arg_name = f'--{name.replace("_", "-")}' + if name == 'dynamic_image_size' or name == 'dynamic_dims': + parser.add_argument( + arg_name, action=DynamicDimsAction, nargs='+', help=desc) + else: + return super().add_argument(parser, name, dtype, default, desc) - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import AscendWrapper - return AscendWrapper(model=backend_files[0], device=device) + +_BackendParam = AscendParam + + +@BACKEND_MANAGERS.register('ascend', param=_BackendParam, ir_param=ONNXParam) +class AscendManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -57,35 +123,155 @@ def get_version(cls) -> str: return 'None' @classmethod - def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + def to_backend(cls, onnx_model: str, output_path: str, + atc_param: AtcParam) -> Sequence[str]: """Convert intermediate representation to given backend. + Example: + >>> from mmdeploy.backend.ascend.onnx2ascend import AtcParam + >>> onnx_path = 'work_dir/end2end.onnx' + >>> output_path = 'work_dir/end2end.om + >>> atc_param = AtcParam(input_shapes=dict(input=[1, 3, 224, 224])) + >>> to_backend(onnx_path, output_path, atc_param) + Args: + onnx_path (ModelProto|str): The path of the onnx model. + output_path (str): Path to save model. + atc_param (AtcParam): The input args to the atc tools. + """ + from .onnx2ascend import from_onnx + from_onnx(onnx_model, output_path, atc_param) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = param.get_model_files() + + input_shapes = param.input_shapes + + atc_param = AtcParam( + input_shapes=input_shapes, + dynamic_batch_size=param.dynamic_batch_size, + dynamic_image_size=param.dynamic_image_size, + dynamic_dims=param.dynamic_dims) + + cls.to_backend(ir_model, model_path, atc_param=atc_param) + + @classmethod + def build_wrapper(cls, model_path: str, device: str = 'cpu'): + """Build the wrapper for the backend model. + + Args: + model_path (str): The om model path. + device (str, optional): The device info. Defaults to 'cpu'. + """ + from .wrapper import AscendWrapper + return AscendWrapper(model=model_path, device=device) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_path = param.get_model_files() + device = param.device + return cls.build_wrapper(model_path, device=device) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + Returns: - Sequence[str]: Backend files. + BaseBackendParam: The packed backend parameter. """ from mmdeploy.utils import get_model_inputs - from .onnx2ascend import from_onnx - model_inputs = get_model_inputs(deploy_cfg) + model_inputs = get_model_inputs(config)[0] + assert 'input_shapes' in model_inputs, ( + 'Can not find model_inputs in config.') + input_shapes = model_inputs['input_shapes'] + dynamic_batch_size = model_inputs.get('dynamic_batch_size', None) + dynamic_image_size = model_inputs.get('dynamic_image_size', None) + dynamic_dims = model_inputs.get('dynamic_dims', None) + + kwargs.setdefault('work_dir', work_dir) + kwargs.setdefault('input_shapes', input_shapes) + kwargs.setdefault('dynamic_batch_size', dynamic_batch_size) + kwargs.setdefault('dynamic_image_size', dynamic_image_size) + kwargs.setdefault('dynamic_dims', dynamic_dims) + + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Import custom modules.') + + parsed_args = parser.parse_args(args) + yield parsed_args + + import_custom_modules(parsed_args.custom_modules) + + # perform command + command = parsed_args._command - om_files = [] - for model_id, onnx_path in enumerate(ir_files): - om_path = osp.splitext(onnx_path)[0] + '.om' - from_onnx(onnx_path, work_dir, model_inputs[model_id]) - om_files.append(om_path) - backend_files = om_files + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + input_shapes=parsed_args.input_shapes, + dynamic_batch_size=parsed_args.dynamic_batch_size, + dynamic_image_size=parsed_args.dynamic_image_size, + dynamic_dims=parsed_args.dynamic_dims, + device=parsed_args.device) - return backend_files + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/ascend/onnx2ascend.py b/mmdeploy/backend/ascend/onnx2ascend.py index a16bb45ebf..425e86b54d 100644 --- a/mmdeploy/backend/ascend/onnx2ascend.py +++ b/mmdeploy/backend/ascend/onnx2ascend.py @@ -1,10 +1,10 @@ # Copyright (c) OpenMMLab. All rights reserved. import os.path as osp import tempfile +from collections import OrderedDict +from dataclasses import dataclass, field from subprocess import call -from typing import Dict, Sequence, Union - -import onnx +from typing import Dict, Sequence from mmdeploy.utils import get_root_logger @@ -17,23 +17,48 @@ def _concat(dims: Sequence) -> str: return ';'.join([','.join(map(str, x)) for x in dims]) -def from_onnx(onnx_model: Union[onnx.ModelProto, str], work_dir: str, - model_inputs: Dict): +@dataclass +class AtcParam: + input_shapes: Dict[str, Sequence] = field(default_factory=OrderedDict) + dynamic_batch_size: Sequence[int] = None + dynamic_image_size: Sequence[Sequence[int]] = None + dynamic_dims: Sequence[Sequence[int]] = None + + def check(self): + dynamic_count = 0 + if self.dynamic_batch_size is not None: + dynamic_count += 1 + if self.dynamic_image_size is not None: + dynamic_count += 1 + if self.dynamic_dims is not None: + dynamic_count += 1 + + if dynamic_count > 1: + raise ValueError('Expect one dynamic flag, but got: ' + f'dynamic_batch_size {self.dynamic_batch_size}; ' + f'dynamic_image_size {self.dynamic_image_size}; ' + f'dynamic_dims {self.dynamic_dims}; ') + + +def from_onnx(onnx_model: str, output_path: str, atc_param: AtcParam): """Convert ONNX to Ascend model. Example: - >>> from mmdeploy.apis.ascend import from_onnx + >>> from mmdeploy.backend.ascend.onnx2ascend import AtcParam, from_onnx >>> onnx_path = 'work_dir/end2end.onnx' - >>> model_inputs = mmengine.Config( - >>> dict(input_shapes=dict(input=[1, 3, 224, 224]))) - >>> from_onnx(onnx_path, work_dir, model_inputs) + >>> output_path = 'work_dir/end2end.om + >>> atc_param = AtcParam(input_shapes=dict(input=[1, 3, 224, 224])) + >>> from_onnx(onnx_path, output_path, atc_param) Args: onnx_path (ModelProto|str): The path of the onnx model. - work_dir (str): Path to load onnx and save model. - model_inputs (Dict): The input args to the atc tools. + output_path (str): Path to save model. + atc_param (AtcParam): The input args to the atc tools. """ + import onnx logger = get_root_logger() + atc_param.check() + if not isinstance(onnx_model, str): onnx_path = tempfile.NamedTemporaryFile(suffix='.onnx').name onnx.save(onnx_model, onnx_path) @@ -41,6 +66,7 @@ def from_onnx(onnx_model: Union[onnx.ModelProto, str], work_dir: str, onnx_path = onnx_model onnx_model = onnx.load(onnx_path) + input_names = [i.name for i in onnx_model.graph.input] for n in onnx_model.graph.node: if n.domain != '': n.domain = '' @@ -48,15 +74,16 @@ def from_onnx(onnx_model: Union[onnx.ModelProto, str], work_dir: str, onnx_model.opset_import.pop(i) onnx.save(onnx_model, onnx_path) - output_path = osp.join(work_dir, osp.splitext(osp.split(onnx_path)[1])[0]) - input_shapes = [] - for name, dims in model_inputs['input_shapes'].items(): + for name, dims in atc_param.input_shapes.items(): input_shapes.append(make_shape_string(name, dims)) input_shapes = ';'.join(input_shapes) - input_format = 'ND' if 'dynamic_dims' in model_inputs else 'NCHW' + input_format = 'ND' if atc_param.dynamic_dims is not None else 'NCHW' + + if output_path.endswith('.om'): + output_path = osp.splitext(output_path)[0] args = [ f'--model={onnx_path}', '--framework=5', f'--output={output_path}', @@ -64,15 +91,26 @@ def from_onnx(onnx_model: Union[onnx.ModelProto, str], work_dir: str, f'--input_shape={input_shapes}' ] - if 'dynamic_batch_size' in model_inputs: - dynamic_batch_size = ','.join( - map(str, model_inputs['dynamic_batch_size'])) + if atc_param.dynamic_batch_size is not None: + dynamic_batch_size = ','.join(map(str, atc_param.dynamic_batch_size)) args.append(f'--dynamic_batch_size={dynamic_batch_size}') - elif 'dynamic_image_size' in model_inputs: - dynamic_image_size = _concat(model_inputs['dynamic_image_size']) + elif atc_param.dynamic_image_size is not None: + dynamic_image_size = atc_param.dynamic_image_size + if isinstance(dynamic_image_size, Dict): + dynamic_image_size = [ + dynamic_batch_size[name] for name in input_names + if name in dynamic_batch_size + ] + dynamic_image_size = _concat(dynamic_image_size) args.append(f'--dynamic_image_size={dynamic_image_size}') - elif 'dynamic_dims' in model_inputs: - dynamic_dims = _concat(model_inputs['dynamic_dims']) + elif atc_param.dynamic_dims is not None: + dynamic_dims = atc_param.dynamic_dims + if isinstance(dynamic_dims, Dict): + dynamic_dims = [ + dynamic_dims[name] for name in input_names + if name in dynamic_dims + ] + dynamic_dims = _concat(dynamic_dims) args.append(f'--dynamic_dims={dynamic_dims}') logger.info(' '.join(('atc', *args))) diff --git a/mmdeploy/backend/ascend/wrapper.py b/mmdeploy/backend/ascend/wrapper.py index 6264a1af85..8401a1ee03 100644 --- a/mmdeploy/backend/ascend/wrapper.py +++ b/mmdeploy/backend/ascend/wrapper.py @@ -238,6 +238,7 @@ def __init__(self): ret = acl.init() if ret == 0: Context.owned_acl = True + Context.ref_count += 1 # add one to prevent early release elif ret == 100002: # ACL_ERROR_REPEAT_INITIALIZE pass else: @@ -306,7 +307,7 @@ class AscendWrapper(BaseWrapper): model (str): Path of the model file. Examples: - >>> from mmdeploy.backend.ascend import AscendWrapper + >>> from mmdeploy.backend.ascend.wrapper import AscendWrapper >>> import torch >>> >>> model_file = 'model.om' @@ -394,7 +395,7 @@ def forward(self, inputs: Dict[str, for binding in self._model_desc.outputs: self._copy_buffer_to_tensor( - self._output.buffers[binding.index], tensor) + self._output.buffers[binding.index], outputs[binding.name]) return outputs diff --git a/mmdeploy/backend/base/__init__.py b/mmdeploy/backend/base/__init__.py index 840c11391b..e0b45ee4ee 100644 --- a/mmdeploy/backend/base/__init__.py +++ b/mmdeploy/backend/base/__init__.py @@ -1,12 +1,21 @@ # Copyright (c) OpenMMLab. All rights reserved. from .backend_manager import (BACKEND_MANAGERS, BaseBackendManager, + BaseBackendParam, FileNameDescriptor, get_backend_manager) from .backend_wrapper_registry import (BACKEND_WRAPPER, get_backend_file_count, get_backend_wrapper_class) from .base_wrapper import BaseWrapper +from .utils import (create_h5pydata_generator, get_obj_by_qualname, + import_custom_modules) __all__ = [ - 'BACKEND_MANAGERS', 'BaseBackendManager', 'get_backend_manager', + 'BACKEND_MANAGERS', 'BaseBackendManager', 'BaseBackendParam', + 'get_backend_manager', 'FileNameDescriptor' +] +__all__ += [ 'BaseWrapper', 'BACKEND_WRAPPER', 'get_backend_wrapper_class', 'get_backend_file_count' ] +__all__ += [ + 'create_h5pydata_generator', 'get_obj_by_qualname', 'import_custom_modules' +] diff --git a/mmdeploy/backend/base/backend_manager.py b/mmdeploy/backend/base/backend_manager.py index 28546ab971..b92020e4e0 100644 --- a/mmdeploy/backend/base/backend_manager.py +++ b/mmdeploy/backend/base/backend_manager.py @@ -1,35 +1,346 @@ # Copyright (c) OpenMMLab. All rights reserved. +import contextlib import importlib -import logging +import os.path as osp +import re from abc import ABCMeta -from typing import Any, Callable, Optional, Sequence +from argparse import Action, ArgumentParser +from collections import OrderedDict +from dataclasses import MISSING, dataclass, field, fields +from typing import (Any, Callable, Dict, Iterable, List, Optional, Sequence, + Union) +from mmdeploy.utils.docstring_parser import inspect_docstring_arguments -class BaseBackendManager(metaclass=ABCMeta): - """Abstract interface of backend manager.""" + +def _parse_shape_type(args: Union[str, List[str]], + allow_placeholder: bool = True) -> Dict[str, List]: + """parse the shape of the arguments. + + Args: + args (Union[str, List[str]]): The arguments string to be parsed. + allow_placeholder (bool, optional): Allow non-int shape. + + Returns: + Dict[str, List]: The parsed shapes + """ + if isinstance(args, str): + args = [args] + + if allow_placeholder: + pattern = r'^[^\S\n]*(((?P\w+):)?(?P(\?|\d+)(x(\?|\d+))*))+$' # noqa + else: + pattern = r'^[^\S\n]*(((?P\w+):)?(?P\d+(x\d+)*))+$' # noqa + + ret = dict() + for arg in args: + arg = [i.strip() for i in arg.split(',')] + for single_arg in arg: + if len(single_arg) == 0: + continue + m = re.match(pattern, single_arg) + if m is None: + raise ValueError(f'Can not parse shape: {single_arg}') + input_name = m.group('input_name') + shape = [ + None if i == '?' else int(i) + for i in m.group('shape').split('x') + ] + if input_name in ret: + raise NameError(f'shape of `{input_name}` has been assigned' + 'more than once.') + ret[input_name] = shape + + return ret + + +class ShapeTypeAction(Action): + """Shape type argparse action.""" + + def __call__(self, parser, namespace, values, option_string=None): + """call action.""" + ret = _parse_shape_type(values, allow_placeholder=False) + setattr(namespace, self.dest, ret) + + +class ShapePlaceHolderAction(Action): + """Shape type argparse action with question mark placeholder.""" + + def __call__(self, parser, namespace, values, option_string=None): + """call action.""" + ret = _parse_shape_type(values, allow_placeholder=True) + setattr(namespace, self.dest, ret) + + +ShapeType = Dict[str, Sequence[int]] + + +class FileNameDescriptor: + """File name descriptor.""" + + def __init__(self, *, default, postfix: str = '', base_name=None): + self._default = default + self._postfix = postfix + self._base_name = base_name + + def __set_name__(self, owner, name): + """set obj name.""" + self._name = '_' + name + + def __get__(self, obj, type): + """file name getter.""" + if obj is None: + return self._default + + # get . + ret = getattr(obj, self._name, self._default) + + # if . is None, try get name from base name + if ret is None and self._base_name is not None: + base_val = getattr(obj, self._base_name, None) + if base_val is not None: + name = osp.splitext(base_val)[0] + ret = name + self._postfix + return ret + + def __set__(self, obj, val): + """file name setter.""" + if val is not None and osp.splitext(val)[1] == '': + val = val + self._postfix + setattr(obj, self._name, val) + + +@dataclass +class BaseBackendParam: + """Base backend parameters. + + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + input_shapes (ShapeType): The Default shape of the inputs. + min_shapes (ShapeType): The minimal shape of the inputs. + max_shapes (ShapeType): The maximal shape of the inputs. + input_names (List[str]): Names of the inputs. + output_names (List[str]): Names of the outputs. + device (str): Device used to perform inference. + quanti_data (Union[Iterable, str]): Iterable object to provide the + quantization data. Each iteration gives a dict of input name and + correspond tensor. + uri (str): The uri of remote device. + """ + _manager = None + + work_dir: str = None + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='') + min_shapes: ShapeType = field(default_factory=OrderedDict) + input_shapes: ShapeType = field(default_factory=OrderedDict) + max_shapes: ShapeType = field(default_factory=OrderedDict) + input_names: List[str] = field(default_factory=list) + output_names: List[str] = field(default_factory=list) + device: str = 'cpu' + quanti_data: Union[Iterable, str] = None + uri: str = None + + def get_model_files(self) -> Union[str, List[str]]: + """get model files.""" + raise NotImplementedError( + f'get_model_files has not implemented for {type(self).__name__}') + + def fix_param(self): + """Fix shapes and names in the parameter.""" + + def _fix_none_name(shapes): + if shapes is not None and None in shapes: + if len(shapes) != 1 or len( + self.input_names) != 1 or self.input_names[0] is None: + raise ValueError( + f'Can not inference name with shapes: {shapes}' + f' and input_names: {self.input_names}') + return {self.input_names[0]: shapes[None]} + else: + return shapes + + def _fill_shape_placeholder(default_shape, shape): + if len(shape) != len(default_shape): + raise ValueError( + f'Can not fill placeholder {shape} with {default_shape}') + return [ + a if a is not None else b + for a, b in zip(shape, default_shape) + ] + + if not isinstance(self.input_shapes, Dict): + raise TypeError('Expect dict input shapes,' + f' but got {type(self.input_shapes)}') + + # fill input names + if len(self.input_names) == 0: + self.input_names = list(self.input_shapes.keys()) + + if len(self.input_names) != len(set(self.input_names)): + raise ValueError( + f'Duplicate names in input_names: {self.input_names}') + + if None in self.input_names: + raise ValueError('Found None in input names.') + + # fix input shapes with no names + self.input_shapes = _fix_none_name(self.input_shapes) + + if self.min_shapes is None or len(self.min_shapes) == 0: + self.min_shapes = self.input_shapes + if self.max_shapes is None or len(self.max_shapes) == 0: + self.max_shapes = self.input_shapes + + if not isinstance(self.min_shapes, Dict): + raise TypeError( + f'Expect min shapes type Dict, got {type(self.min_shapes)}.') + if not isinstance(self.max_shapes, Dict): + raise TypeError( + f'Expect max shapes type Dict, got {type(self.max_shapes)}.') + + self.min_shapes = _fix_none_name(self.min_shapes) + self.max_shapes = _fix_none_name(self.max_shapes) + + # fix placeholder min/max shapes + for name, in_shape in self.input_shapes.items(): + if name in self.min_shapes: + self.min_shapes[name] = _fill_shape_placeholder( + in_shape, self.min_shapes[name]) + + if name in self.max_shapes: + self.max_shapes[name] = _fill_shape_placeholder( + in_shape, self.max_shapes[name]) + + def check_param(self): + """Check the parameter validation.""" + self.fix_param() + + input_shapes = self.input_shapes + min_shapes = self.min_shapes + max_shapes = self.max_shapes + + if not (len(input_shapes) == len(min_shapes) == len(max_shapes)): + raise ValueError(f'len(min_shapes) = {len(min_shapes)}\n', + f'len(input_shapes) = {len(input_shapes)}\n', + f'len(max_shapes) = {len(max_shapes)}\n', + ' should be the same.') + + for name, input_shape in input_shapes.items(): + if name not in min_shapes: + raise NameError(f'{name} not found in min_shapes.') + if name not in max_shapes: + raise NameError(f'{name} not found in max_shapes.') + min_shape = min_shapes[name] + max_shape = max_shapes[name] + + if not isinstance(input_shape, Sequence): + raise TypeError(f'input shape of {name} is not sequence.') + if not isinstance(min_shape, Sequence): + raise TypeError(f'min shape of {name} is not sequence.') + if not isinstance(max_shape, Sequence): + raise TypeError(f'max shape of {name} is not sequence.') + if not (len(input_shape) == len(min_shape) == len(max_shape)): + raise ValueError( + f'len(min_shapes[{name}]) = {len(min_shape)}\n', + f'len(input_shapes[{name}]) = {len(input_shape)}\n', + f'len(max_shapes[{name}]) = {len(max_shape)}\n', + ' should be the same.') + + for min_s, opt_s, max_s in zip(min_shape, input_shape, max_shape): + if not min_s <= opt_s <= max_s: + raise ValueError( + f'Input {name} has invalid shape:\n', + f'min shape:{min_shape}\n', + f'input shape:{input_shape}\n', + f'max shape:{max_shape}', + ) @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + def get_manager(cls): + """Get backend manager.""" + return cls._manager + + @classmethod + def add_argument(cls, parser: ArgumentParser, name: str, dtype: Any, + default: Any, desc: str): + """Add argument to the parser. Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. + parser (ArgumentParser): Parser object to add argument. + name (str): Name of the argument. + dtype (Any): Argument type. + default (Any): Default value of the argument. + desc (str): Description of the argument. """ - raise NotImplementedError( - f'build_wrapper has not been implemented for `{cls.__name__}`') + arg_name = f'--{name.replace("_", "-")}' + if dtype == bool: + if default is True: + action = 'store_false' + else: + action = 'store_true' + parser.add_argument(arg_name, action=action, help=desc) + elif dtype == FileNameDescriptor: + parser.add_argument(arg_name, type=str, default=default, help=desc) + elif dtype == ShapeType: + action = ShapeTypeAction \ + if name == 'input_shapes' else ShapePlaceHolderAction + parser.add_argument( + arg_name, action=action, nargs='+', default=default, help=desc) + elif dtype == List[str]: + parser.add_argument( + arg_name, type=str, nargs='+', default=default, help=desc) + else: + parser.add_argument( + arg_name, type=dtype, default=default, help=desc) + + @classmethod + def add_arguments(cls, + parser: ArgumentParser, + ignore_fields: Optional[List] = None): + """Add Arguments to the parser to build the param. + + Args: + parser (ArgumentParser): Parser to add arguments. + ignore_fields (Optional[List], optional): Ignore some fields in + the dataclass. Defaults to None. + """ + parser.description = f'build {cls.__name__}' + + if ignore_fields is None: + ignore_fields = [] + remain_fields = inspect_docstring_arguments( + cls, ignore_args=ignore_fields) + + field_map = dict((f.name, f) for f in fields(cls)) + + for remain_field in remain_fields: + name = remain_field.name + desc = remain_field.desc + + assert name in field_map, \ + f'{name} is not a field in {cls.__name__}' + cls_field = field_map[name] + + dtype = cls_field.type + if cls_field.default is not MISSING: + default = cls_field.default + elif isinstance(cls_field.default_factory, Callable): + default = cls_field.default_factory() + else: + default = None + + cls.add_argument( + parser, name, dtype=dtype, default=default, desc=desc) + + +class BaseBackendManager(metaclass=ABCMeta): + """Abstract interface of backend manager.""" + + build_ir_param = None + build_param = BaseBackendParam @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -75,28 +386,88 @@ def check_env(cls, log_callback: Callable = lambda _: _) -> str: return info @classmethod - def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + def to_backend(cls, ir_path, *args, **kwargs): """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. - Returns: - Sequence[str]: Backend files. + ir_path (str): The intermediate representation files. """ raise NotImplementedError( f'to_backend has not been implemented for `{cls.__name__}`') + @classmethod + def to_backend_from_param(cls, ir_model: str, param: BaseBackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + raise NotImplementedError( + 'to_backend_from_param has not been implemented for ' + f'`{cls.__name__}`') + + @classmethod + def build_wrapper(cls, *args, **kwargs): + """Build the wrapper for the backend model.""" + raise NotImplementedError( + f'build_wrapper has not been implemented for `{cls.__name__}`') + + @classmethod + def build_wrapper_from_param(cls, param: BaseBackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + raise NotImplementedError( + 'build_wrapper_from_param has not been implemented for ' + f'`{cls.__name__}`') + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + raise NotImplementedError( + f'parse_args of {cls.__name__} has not been implemented.') + + @classmethod + def main(cls): + """Create console tools.""" + parser = ArgumentParser() + with cls.parse_args(parser): + pass + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: List[str] = None, + **kwargs) -> BaseBackendParam: + """Build param from deploy config. + + This is a bridge between old and new api. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + raise NotImplementedError( + 'build_param_from_config has not been implemented' + f' for {cls.__name__}.') + class BackendManagerRegistry: """backend manager registry.""" @@ -104,7 +475,11 @@ class BackendManagerRegistry: def __init__(self): self._module_dict = {} - def register(self, name: str, enum_name: Optional[str] = None): + def register(self, + name: str, + enum_name: Optional[str] = None, + param: Any = None, + ir_param: Any = None): """register backend manager. Args: @@ -135,6 +510,10 @@ def wrap_manager(cls): self._module_dict[name] = cls cls.backend_name = name + cls.build_param = param + cls.build_ir_param = ir_param + if param is not None: + param._manager = cls return cls diff --git a/mmdeploy/backend/base/utils.py b/mmdeploy/backend/base/utils.py new file mode 100644 index 0000000000..8e4b1b6ef9 --- /dev/null +++ b/mmdeploy/backend/base/utils.py @@ -0,0 +1,97 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import importlib +from typing import Any, Dict, Sequence, Union + + +def get_obj_by_qualname(qualname: str) -> Any: + """Get object by the qualname. + + Args: + qualname (str): The qualname of the object + + Returns: + Any: The object with qualname + """ + split_qualname = qualname.split('.') + for i in range(len(split_qualname), 0, -1): + try: + exec('import {}'.format('.'.join(split_qualname[:i]))) + break + except Exception: + continue + + obj = eval(qualname) + + return obj + + +def create_h5pydata_generator(data_file: Union[str, Any], + input_shapes: Dict[str, Sequence[int]], + data_type: str = 'end2end'): + """Create data generator for h5py data. + + Args: + data_file (Union[str, Any]): h5py file. + input_shapes (Dict[str, Sequence]): Input shape of each input tensors. + data_type (str, optional): Data type id. Defaults to 'end2end'. + """ + import h5py + import numpy as np + if isinstance(data_file, str): + data_file = h5py.File(data_file, mode='r') + + try: + assert 'calib_data' in data_file + calib_data = data_file['calib_data'] + assert data_type in calib_data + calib_data = calib_data[data_type] + + names = list(calib_data.keys()) + first_input_group = calib_data[list(calib_data.keys())[0]] + dataset_length = len(first_input_group) + + # iterate over all data + for idx in range(dataset_length): + + yield_data = dict() + for name in names: + input_group = calib_data[name] + data_np = input_group[str(idx)][...] + + # tile the tensor so we can keep the same distribute + opt_shape = input_shapes[name] + data_shape = data_np.shape + + reps = [ + int(np.ceil(opt_s / data_s)) + for opt_s, data_s in zip(opt_shape, data_shape) + ] + + data_np = np.tile(data_np, reps) + + slice_list = tuple(slice(0, end) for end in opt_shape) + data_np = data_np[slice_list] + + yield_data[name] = data_np + + yield yield_data + + except Exception as e: + raise e + finally: + data_file.close() + + +def import_custom_modules(custom_modules: Sequence): + """Import custom module.""" + from mmdeploy.utils import get_root_logger + logger = get_root_logger(0) + custom_modules = [] if custom_modules is None else custom_modules + + for qualname in custom_modules: + try: + importlib.import_module(qualname) + logger.info(f'Import custom module: {qualname}') + except Exception as e: + logger.warning('Failed to import custom module: ' + f'{qualname} with error: {e}') diff --git a/mmdeploy/backend/coreml/__init__.py b/mmdeploy/backend/coreml/__init__.py index b44447de9f..6d99d8401e 100644 --- a/mmdeploy/backend/coreml/__init__.py +++ b/mmdeploy/backend/coreml/__init__.py @@ -1,15 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import CoreMLManager +from .backend_manager import CoreMLManager, CoreMLParam _BackendManager = CoreMLManager - is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['CoreMLManager'] - -if is_available(): - from . import ops - from .torchscript2coreml import get_model_suffix - from .wrapper import CoreMLWrapper - __all__ += ['CoreMLWrapper', 'get_model_suffix', 'ops'] +__all__ = ['CoreMLParam', 'CoreMLManager'] diff --git a/mmdeploy/backend/coreml/backend_manager.py b/mmdeploy/backend/coreml/backend_manager.py index f81316e00b..c444869830 100644 --- a/mmdeploy/backend/coreml/backend_manager.py +++ b/mmdeploy/backend/coreml/backend_manager.py @@ -1,36 +1,73 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os.path as osp -from typing import Any, Optional, Sequence +from argparse import ArgumentParser +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.torchscript import TorchScriptParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor, import_custom_modules) -@BACKEND_MANAGERS.register('coreml') -class CoreMLManager(BaseBackendManager): +@dataclass +class CoreMLParam(BaseBackendParam): + """CoreML backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + input_shapes (ShapeType): The Default shape of the inputs. + min_shapes (ShapeType): The minimal shape of the inputs. + max_shapes (ShapeType): The maximal shape of the inputs. + input_names (List[str]): Names of the inputs. + output_names (List[str]): Names of the outputs. + compute_precision (str): The model precision, FLOAT16 or FLOAT32, + read coremltools.precision for more detail, default `FLOAT32`. + convert_to (str): The converted model type, can be + 'neuralnetwork' or 'mlprogram'. Defaults to 'neuralnetwork'. + minimum_deployment_target (str, optional): minimum deploy target. + iOS15, iOS16, etc., see coremltools.target + skip_model_load (bool, optional): Skip model load. + Defaults to True. + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.mlpackage') + compute_precision: str = 'FLOAT32' + convert_to: str = 'mlprogram' + minimum_deployment_target: Optional[str] = None + skip_model_load: bool = True - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import CoreMLWrapper - return CoreMLWrapper(model_file=backend_files[0]) + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str), ('Expect string work_dir, ' + f'got {self.work_dir}') + assert isinstance(self.file_name, str), ('Expect string file_name, ' + f'got {self.file_name}') + file_name = self.file_name + return osp.join(self.work_dir, file_name) + + def check_param(self): + """Check the parameter validation.""" + if self.convert_to == 'mlprogram' and not self.file_name.endswith( + '.mlpackage'): + raise ValueError('extension should be `.mlpackage` when ' + 'convert_to == `mlprogram`. ') + if self.convert_to == 'neuralnetwork' and not self.file_name.endswith( + '.mlmodel'): + raise ValueError('extension should be `.mlmodel` when ' + 'convert_to == `neuralnetwork`. ') + + super().check_param() + + +_BackendParam = CoreMLParam + + +@BACKEND_MANAGERS.register( + 'coreml', param=_BackendParam, ir_param=TorchScriptParam) +class CoreMLManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -58,66 +95,220 @@ def get_version(cls) -> str: @classmethod def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + torchscript_model: str, + output_path: str, + input_names: Sequence[str], + output_names: Sequence[str], + input_shapes: Dict[str, Dict], + min_shapes: Dict[str, Dict] = None, + max_shapes: Dict[str, Dict] = None, + compute_precision: str = 'FLOAT32', + convert_to: str = None, + minimum_deployment_target: Optional[str] = None, + skip_model_load: bool = True) -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. + torchscript_model (Union[str, torch.jit.RecursiveScriptModule]): + The torchscript model to be converted. + output_path (str): The output file. + input_names (Sequence[str]): The input names of the model. + output_names (Sequence[str]): The output names of the model. + input_shapes (ShapeType): The Default shape of the inputs. + min_shapes (ShapeType): The minimal shape of the inputs. + max_shapes (ShapeType): The maximal shape of the inputs. + compute_precision (str): The model precision, FLOAT16 or FLOAT32, + read coremltools.precision for more detail, default `FLOAT32`. + convert_to (str): The converted model type, can be + 'neuralnetwork' or 'mlprogram'. Defaults to 'neuralnetwork'. + minimum_deployment_target (str, optional): minimum deploy target. + iOS15, iOS16, etc., see coremltools.target + skip_model_load (bool, optional): Skip model load. + Defaults to True. Returns: Sequence[str]: Backend files. """ + from .torchscript2coreml import from_torchscript + + from_torchscript( + torchscript_model, + output_path, + input_names=input_names, + output_names=output_names, + input_shapes=input_shapes, + min_shapes=min_shapes, + max_shapes=max_shapes, + compute_precision=compute_precision, + convert_to=convert_to, + minimum_deployment_target=minimum_deployment_target, + skip_model_load=skip_model_load) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + param.check_param() + + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = param.get_model_files() + + minimum_deployment_target = param.minimum_deployment_target + cls.to_backend( + ir_model, + model_path, + input_names=param.input_names, + output_names=param.output_names, + input_shapes=param.input_shapes, + min_shapes=param.min_shapes, + max_shapes=param.max_shapes, + compute_precision=param.compute_precision, + convert_to=param.convert_to, + minimum_deployment_target=minimum_deployment_target, + skip_model_load=param.skip_model_load) + + @classmethod + def build_wrapper(cls, model_path: str): + """Build the wrapper for the backend model. + + Args: + model_path (str): Backend files. + """ + from .wrapper import CoreMLWrapper + return CoreMLWrapper(model_file=model_path) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_path = param.get_model_files() + return cls.build_wrapper(model_path) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + from mmdeploy.utils import (get_common_config, get_ir_config, get_model_inputs, load_config) - from .torchscript2coreml import from_torchscript, get_model_suffix - - coreml_files = [] - for model_id, torchscript_path in enumerate(ir_files): - torchscript_name = osp.splitext(osp.split(torchscript_path)[1])[0] - output_file_prefix = osp.join(work_dir, torchscript_name) - - deploy_cfg = load_config(deploy_cfg)[0] - - common_params = get_common_config(deploy_cfg) - model_params = get_model_inputs(deploy_cfg)[model_id] - - final_params = common_params - final_params.update(model_params) - - ir_config = get_ir_config(deploy_cfg) - input_names = ir_config.get('input_names', []) - output_names = ir_config.get('output_names', []) - input_shapes = final_params['input_shapes'] - compute_precision = final_params.get('compute_precision', - 'FLOAT32') - convert_to = deploy_cfg.backend_config.convert_to - - minimum_deployment_target = final_params.get( - 'minimum_deployment_target', None) - skip_model_load = final_params.get('skip_model_load', False) - - from_torchscript( - torchscript_path, - output_file_prefix, - input_names=input_names, - output_names=output_names, - input_shapes=input_shapes, - compute_precision=compute_precision, - convert_to=convert_to, - minimum_deployment_target=minimum_deployment_target, - skip_model_load=skip_model_load) - - suffix = get_model_suffix(convert_to) - output_path = output_file_prefix + suffix - coreml_files.append(output_path) - - return coreml_files + + deploy_cfg = config + deploy_cfg = load_config(deploy_cfg)[0] + + common_params = get_common_config(deploy_cfg) + model_params = get_model_inputs(deploy_cfg)[0] + + final_params = common_params + final_params.update(model_params) + + ir_config = get_ir_config(deploy_cfg) + input_names = ir_config.get('input_names', []) + output_names = ir_config.get('output_names', []) + input_shapes = final_params['input_shapes'] + min_shapes = dict( + (name, shape['min_shape']) for name, shape in input_shapes.items()) + max_shapes = dict( + (name, shape['max_shape']) for name, shape in input_shapes.items()) + input_shapes = dict((name, shape['default_shape']) + for name, shape in input_shapes.items()) + compute_precision = final_params.get('compute_precision', 'FLOAT32') + convert_to = deploy_cfg.backend_config.convert_to + + minimum_deployment_target = final_params.get( + 'minimum_deployment_target', None) + skip_model_load = final_params.get('skip_model_load', False) + + kwargs.setdefault('work_dir', work_dir) + kwargs.setdefault('input_shapes', input_shapes) + kwargs.setdefault('min_shapes', min_shapes) + kwargs.setdefault('max_shapes', max_shapes) + kwargs.setdefault('input_names', input_names) + kwargs.setdefault('output_names', output_names) + kwargs.setdefault('compute_precision', compute_precision) + kwargs.setdefault('convert_to', convert_to) + kwargs.setdefault('minimum_deployment_target', + minimum_deployment_target) + kwargs.setdefault('skip_model_load', skip_model_load) + + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from torchscript model.') + export_parser.add_argument( + '--torchscript-path', + required=True, + help='torchscript model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Import custom modules.') + + parsed_args = parser.parse_args(args) + yield parsed_args + + import_custom_modules(parsed_args.custom_modules) + # perform command + command = parsed_args._command + + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + input_names=parsed_args.input_names, + output_names=parsed_args.output_names, + input_shapes=parsed_args.input_shapes, + min_shapes=parsed_args.min_shapes, + max_shapes=parsed_args.max_shapes, + compute_precision=parsed_args.compute_precision, + convert_to=parsed_args.convert_to, + minimum_deployment_target=parsed_args. + minimum_deployment_target, + skip_model_load=parsed_args.skip_model_load) + + cls.to_backend_from_param(parsed_args.torchscript_path, param) diff --git a/mmdeploy/backend/coreml/torchscript2coreml.py b/mmdeploy/backend/coreml/torchscript2coreml.py index 6d2b52a7d3..61bdb32fda 100644 --- a/mmdeploy/backend/coreml/torchscript2coreml.py +++ b/mmdeploy/backend/coreml/torchscript2coreml.py @@ -1,11 +1,12 @@ # Copyright (c) OpenMMLab. All rights reserved. - +import os.path as osp from typing import Dict, Optional, Sequence, Union import coremltools as ct import torch from mmdeploy.utils import get_root_logger +from . import ops # noqa try: # user might need ops from torchvision @@ -13,6 +14,8 @@ except ImportError: pass +SUFFIX_MODE_MAP = {'.mlmodel': 'neuralnetwork', '.mlpackage': 'mlprogram'} + def get_model_suffix(convert_to: str) -> str: assert convert_to == 'neuralnetwork' or convert_to == 'mlprogram' @@ -24,11 +27,9 @@ def get_model_suffix(convert_to: str) -> str: return suffix -def create_shape(name: str, input_shapes: Dict) -> ct.Shape: +def create_shape(name: str, default_shape: Sequence, min_shape: Sequence, + max_shape: Sequence) -> ct.Shape: """Create input shape.""" - min_shape = input_shapes['min_shape'] - max_shape = input_shapes['max_shape'] - default_shape = input_shapes['default_shape'] assert len(min_shape) == len(max_shape) == len(default_shape) shape = [] n_dim = len(min_shape) @@ -49,25 +50,27 @@ def create_shape(name: str, input_shapes: Dict) -> ct.Shape: def from_torchscript(torchscript_model: Union[str, torch.jit.RecursiveScriptModule], - output_file_prefix: str, + output_path: str, input_names: Sequence[str], output_names: Sequence[str], input_shapes: Dict[str, Dict], + min_shapes: Dict[str, Dict] = None, + max_shapes: Dict[str, Dict] = None, compute_precision: str = 'FLOAT32', - convert_to: str = 'neuralnetwork', + convert_to: str = None, minimum_deployment_target: Optional[str] = None, - skip_model_load: bool = True, - **kwargs): + skip_model_load: bool = True): """Create a coreml engine from torchscript. Args: torchscript_model (Union[str, torch.jit.RecursiveScriptModule]): The torchscript model to be converted. - output_file_prefix (str): The output file prefix. + output_path (str): The output file. input_names (Sequence[str]): The input names of the model. output_names (Sequence[str]): The output names of the model. - input_shapes (Dict): The input shapes include max_shape, min_shape and - default_shape + input_shapes (ShapeType): The Default shape of the inputs. + min_shapes (ShapeType): The minimal shape of the inputs. + max_shapes (ShapeType): The maximal shape of the inputs. compute_precision (str): The model precision, FLOAT16 or FLOAT32, see coremltools.precision, default `FLOAT32`. convert_to (str): The converted model type, can be @@ -92,13 +95,28 @@ def from_torchscript(torchscript_model: Union[str, inputs = [] outputs = [] + if min_shapes is None: + min_shapes = input_shapes + if max_shapes is None: + max_shapes = input_shapes for name in input_names: - shape = create_shape(name, input_shapes[name]) + input_shape = input_shapes[name] + min_shape = min_shapes.get(name, input_shape) + max_shape = max_shapes.get(name, input_shape) + shape = create_shape(name, input_shape, min_shape, max_shape) inputs.append(shape) for name in output_names: outputs.append(ct.TensorType(name=name)) + if convert_to is None: + suffix = osp.splitext(output_path)[1] + convert_to = SUFFIX_MODE_MAP[suffix] + + if convert_to not in ['neuralnetwork', 'mlprogram']: + get_root_logger().warning(f'Unknown postfix: {convert_to}. ', + 'Use default mode: neuralnetwork.') + if convert_to == 'neuralnetwork': compute_precision = None else: @@ -114,7 +132,4 @@ def from_torchscript(torchscript_model: Union[str, if minimum_deployment_target else None, skip_model_load=skip_model_load) - suffix = get_model_suffix(convert_to) - output_path = output_file_prefix + suffix - mlmodel.save(output_path) diff --git a/mmdeploy/backend/coreml/wrapper.py b/mmdeploy/backend/coreml/wrapper.py index 1444d1a3d3..5ba925ac38 100644 --- a/mmdeploy/backend/coreml/wrapper.py +++ b/mmdeploy/backend/coreml/wrapper.py @@ -20,7 +20,7 @@ class CoreMLWrapper(BaseWrapper): bin_file (str): Path of a binary file. Examples: - >>> from mmdeploy.backend.coreml import CoreMLWrapper + >>> from mmdeploy.backend.coreml.wrapper import CoreMLWrapper >>> import torch >>> >>> model_file = 'model.mlpackage' diff --git a/mmdeploy/backend/ncnn/__init__.py b/mmdeploy/backend/ncnn/__init__.py index 9e3f65f35f..ad8a2550dd 100644 --- a/mmdeploy/backend/ncnn/__init__.py +++ b/mmdeploy/backend/ncnn/__init__.py @@ -1,18 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import NCNNManager +from .backend_manager import NCNNManager, NCNNParam from .onnx2ncnn import from_onnx _BackendManager = NCNNManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['NCNNManager', 'from_onnx'] - -if is_available(): - try: - from .wrapper import NCNNWrapper - - __all__ += ['NCNNWrapper'] - except Exception: - pass +__all__ = ['NCNNParam', 'NCNNManager', 'from_onnx'] diff --git a/mmdeploy/backend/ncnn/backend_manager.py b/mmdeploy/backend/ncnn/backend_manager.py index 6b8d29b69b..8e078cf4ed 100644 --- a/mmdeploy/backend/ncnn/backend_manager.py +++ b/mmdeploy/backend/ncnn/backend_manager.py @@ -1,50 +1,51 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os.path as osp import sys -from typing import Any, Callable, Optional, Sequence +from argparse import ArgumentParser +from dataclasses import dataclass +from typing import Any, Callable, List, Optional, Sequence -from mmdeploy.utils import get_backend_config, get_root_logger -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor) -@BACKEND_MANAGERS.register('ncnn') -class NCNNManager(BaseBackendManager): +@dataclass +class NCNNParam(BaseBackendParam): + """NCNN backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + bin_name (str): Serialized bin file. If not given, bin_name would be + the same as file_name with postfix `.param` + use_vulkan (str): Perform inference with vulkan. + precision (str): Precision of the model, `INT8` or `FP32` + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.param') + bin_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.bin', base_name='file_name') + use_vulkan: bool = False + precision: str = 'FP32' - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import NCNNWrapper + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str) + assert isinstance(self.file_name, str) + param_file_path = osp.join(self.work_dir, self.file_name) + assert isinstance(self.bin_name, str) + bin_file_path = osp.join(self.work_dir, self.bin_name) + return param_file_path, bin_file_path - # For unittest deploy_config will not pass into _build_wrapper - # function. - if deploy_cfg: - backend_config = get_backend_config(deploy_cfg) - use_vulkan = backend_config.get('use_vulkan', False) - else: - use_vulkan = False - return NCNNWrapper( - param_file=backend_files[0], - bin_file=backend_files[1], - output_names=output_names, - use_vulkan=use_vulkan) + +_BackendParam = NCNNParam + + +@BACKEND_MANAGERS.register('ncnn', param=_BackendParam, ir_param=ONNXParam) +class NCNNManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -103,43 +104,156 @@ def check_env(cls, log_callback: Callable = lambda _: _) -> str: return info @classmethod - def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + def to_backend(cls, onnx_path: str, param_path: str, + bin_path: str) -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. + onnx_path (str): The ONNX model path. + param_path (str): ncnn parameter file path. + bin_path (str): ncnn bin file path. work_dir (str): The work directory, backend files and logs should be saved in this directory. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. Returns: Sequence[str]: Backend files. """ + from mmdeploy.utils import get_root_logger logger = get_root_logger() - from . import is_available - - if not is_available(): + if not cls.is_available(): logger.error('ncnn support is not available, please make sure:\n' '1) `mmdeploy_onnx2ncnn` existed in `PATH`\n' '2) python import ncnn success') sys.exit(1) - from mmdeploy.apis.ncnn import get_output_model_file from .onnx2ncnn import from_onnx - backend_files = [] - for onnx_path in ir_files: - model_param_path, model_bin_path = get_output_model_file( - onnx_path, work_dir) - onnx_name = osp.splitext(osp.split(onnx_path)[1])[0] - from_onnx(onnx_path, osp.join(work_dir, onnx_name)) + from_onnx(onnx_path, param_path, bin_path) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + assert isinstance(param.bin_name, str) + bin_path = osp.join(param.work_dir, param.bin_name) + + cls.to_backend(ir_model, model_path, bin_path) + + @classmethod + def build_wrapper(cls, + param_path: str, + bin_path: str, + output_names: Optional[Sequence[str]] = None, + use_vulkan: bool = False): + """Build the wrapper for the backend model. + + Args: + param_path (str): ncnn parameter file path. + bin_path (str): ncnn bin file path. + device (str, optional): The device info. Defaults to 'cpu'. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + use_vulkan (str): Perform inference with vulkan. + """ + from .wrapper import NCNNWrapper + + # For unittest deploy_config will not pass into _build_wrapper + # function. + return NCNNWrapper( + param_file=param_path, + bin_file=bin_path, + output_names=output_names, + use_vulkan=use_vulkan) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + param_path, bin_path = param.get_model_files() + output_names = param.output_names + if output_names is not None and len(output_names) == 0: + output_names = None + return cls.build_wrapper( + param_path, bin_path, output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + from mmdeploy.utils import get_backend_config + backend_cfg = get_backend_config(config) + use_vulkan = backend_cfg.get('use_vulkan', False) + kwargs.update(dict(work_dir=work_dir, use_vulkan=use_vulkan)) + + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + if len(backend_files) > 1: + kwargs['bin_name'] = backend_files[1] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + + parsed_args = parser.parse_args(args) + yield parsed_args + + # perform command + command = parsed_args._command - backend_files += [model_param_path, model_bin_path] + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + bin_name=parsed_args.bin_name, + precision=parsed_args.precision, + use_vulkan=parsed_args.use_vulkan) - return backend_files + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/ncnn/onnx2ncnn.py b/mmdeploy/backend/ncnn/onnx2ncnn.py index 86f9a8e4f2..a93024db5d 100644 --- a/mmdeploy/backend/ncnn/onnx2ncnn.py +++ b/mmdeploy/backend/ncnn/onnx2ncnn.py @@ -39,8 +39,8 @@ def get_output_model_file(onnx_path: str, return [save_param, save_bin] -def from_onnx(onnx_model: Union[onnx.ModelProto, str], - output_file_prefix: str): +def from_onnx(onnx_model: Union[onnx.ModelProto, str], param_path: str, + bin_path: str): """Convert ONNX to ncnn. The inputs of ncnn include a model file and a weight file. We need to use @@ -50,12 +50,14 @@ def from_onnx(onnx_model: Union[onnx.ModelProto, str], Example: >>> from mmdeploy.apis.ncnn import from_onnx >>> onnx_path = 'work_dir/end2end.onnx' - >>> output_file_prefix = 'work_dir/end2end' - >>> from_onnx(onnx_path, output_file_prefix) + >>> param_path = 'work_dir/end2end.param' + >>> bin_path = 'work_dir/end2end.bin' + >>> from_onnx(onnx_path, param_path, bin_path) Args: onnx_path (ModelProto|str): The path of the onnx model. - output_file_prefix (str): The path to save the output ncnn file. + param_path (str): The path to save the output ncnn param file. + bin_path (str): The path to save the output ncnn bin file. """ if not isinstance(onnx_model, str): @@ -64,9 +66,6 @@ def from_onnx(onnx_model: Union[onnx.ModelProto, str], else: onnx_path = onnx_model - save_param = output_file_prefix + '.param' - save_bin = output_file_prefix + '.bin' - onnx2ncnn_path = get_onnx2ncnn_path() - ret_code = call([onnx2ncnn_path, onnx_path, save_param, save_bin]) + ret_code = call([onnx2ncnn_path, onnx_path, param_path, bin_path]) assert ret_code == 0, 'onnx2ncnn failed' diff --git a/mmdeploy/backend/ncnn/wrapper.py b/mmdeploy/backend/ncnn/wrapper.py index adc9dccdf0..e468030245 100644 --- a/mmdeploy/backend/ncnn/wrapper.py +++ b/mmdeploy/backend/ncnn/wrapper.py @@ -23,7 +23,7 @@ class NCNNWrapper(BaseWrapper): ncnn model. Examples: - >>> from mmdeploy.backend.ncnn import NCNNWrapper + >>> from mmdeploy.backend.ncnn.wrapper import NCNNWrapper >>> import torch >>> >>> param_file = 'model.params' diff --git a/mmdeploy/backend/onnxruntime/__init__.py b/mmdeploy/backend/onnxruntime/__init__.py index 8dd174c2d9..947f48ab1a 100644 --- a/mmdeploy/backend/onnxruntime/__init__.py +++ b/mmdeploy/backend/onnxruntime/__init__.py @@ -1,17 +1,17 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import ONNXRuntimeManager +from .backend_manager import ONNXRuntimeManager, ONNXRuntimeParam _BackendManager = ONNXRuntimeManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['ONNXRuntimeManager'] +__all__ = ['ONNXRuntimeParam', 'ONNXRuntimeManager'] if is_available(): - try: - # import wrapper if pytorch is available - from .wrapper import ORTWrapper - __all__ += ['ORTWrapper'] - except Exception: - pass + from .wrapper import ORTWrapper + + __all__ += ['ORTWrapper'] diff --git a/mmdeploy/backend/onnxruntime/backend_manager.py b/mmdeploy/backend/onnxruntime/backend_manager.py index 7410f381c0..ead2a8ad6f 100644 --- a/mmdeploy/backend/onnxruntime/backend_manager.py +++ b/mmdeploy/backend/onnxruntime/backend_manager.py @@ -1,40 +1,42 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging import os.path as osp -from typing import Any, Callable, Optional, Sequence +import shutil +from dataclasses import dataclass +from typing import Any, Callable, List, Optional, Sequence -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor) -@BACKEND_MANAGERS.register('onnxruntime') -class ONNXRuntimeManager(BaseBackendManager): +@dataclass +class ONNXRuntimeParam(BaseBackendParam): + """ONNX Runtime backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + device (str): The device used to perform the inference. Default to cpu. + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.onnx') - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str), ('Expect string work_dir, ' + f'got {self.work_dir}') + assert isinstance(self.file_name, str), ('Expect string file_name, ' + f'got {self.file_name}') + return osp.join(self.work_dir, self.file_name) + + +_BackendParam = ONNXRuntimeParam - from .wrapper import ORTWrapper - return ORTWrapper( - onnx_file=backend_files[0], - device=device, - output_names=output_names) + +@BACKEND_MANAGERS.register( + 'onnxruntime', param=_BackendParam, ir_param=ONNXParam) +class ONNXRuntimeManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -122,21 +124,81 @@ def check_env(cls, log_callback: Callable = lambda _: _) -> str: return info @classmethod - def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + def to_backend(cls, onnx_path: str, save_path: str): """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. + onnx_path (str): The intermediate representation files. + save_path (str): The save path of onnx path. Returns: Sequence[str]: Backend files. """ - return ir_files + if osp.abspath(save_path) != osp.abspath(onnx_path): + shutil.copy(onnx_path, save_path) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: BaseBackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + save_path = osp.join(param.work_dir, param.file_name) + cls.to_backend(ir_model, save_path) + + @classmethod + def build_wrapper(cls, + model_path: str, + device: str = 'cpu', + output_names: Optional[Sequence[str]] = None): + """Build the wrapper for the backend model. + + Args: + model_path (str): ONNX model file. + device (str, optional): The device info. Defaults to 'cpu'. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + """ + + from .wrapper import ORTWrapper + return ORTWrapper( + onnx_file=model_path, device=device, output_names=output_names) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (_BackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + output_names = param.output_names + if output_names is not None and len(output_names) == 0: + output_names = None + return cls.build_wrapper( + model_path, device=param.device, output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: List[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + return _BackendParam( + work_dir=work_dir, file_name=backend_files[0], **kwargs) diff --git a/mmdeploy/backend/onnxruntime/wrapper.py b/mmdeploy/backend/onnxruntime/wrapper.py index a620ce9139..8e150838a8 100644 --- a/mmdeploy/backend/onnxruntime/wrapper.py +++ b/mmdeploy/backend/onnxruntime/wrapper.py @@ -23,7 +23,7 @@ class ORTWrapper(BaseWrapper): model. Examples: - >>> from mmdeploy.backend.onnxruntime import ORTWrapper + >>> from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper >>> import torch >>> >>> onnx_file = 'model.onnx' diff --git a/mmdeploy/backend/openvino/__init__.py b/mmdeploy/backend/openvino/__init__.py index f17ae36c2e..c0c367e8e5 100644 --- a/mmdeploy/backend/openvino/__init__.py +++ b/mmdeploy/backend/openvino/__init__.py @@ -1,16 +1,12 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import OpenVINOManager +from .backend_manager import OpenVINOManager, OpenVINOParam _BackendManager = OpenVINOManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['OpenVINOManager'] -if is_available(): - from .onnx2openvino import get_output_model_file - from .utils import ModelOptimizerOptions - from .wrapper import OpenVINOWrapper - __all__ += [ - 'OpenVINOWrapper', 'get_output_model_file', 'ModelOptimizerOptions' - ] +__all__ = ['OpenVINOParam', 'OpenVINOManager'] diff --git a/mmdeploy/backend/openvino/backend_manager.py b/mmdeploy/backend/openvino/backend_manager.py index 63ebc92860..e8395aaec6 100644 --- a/mmdeploy/backend/openvino/backend_manager.py +++ b/mmdeploy/backend/openvino/backend_manager.py @@ -1,36 +1,51 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging -from typing import Any, Optional, Sequence +import contextlib +import os.path as osp +import tempfile +from argparse import ArgumentParser +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor) -@BACKEND_MANAGERS.register('openvino') -class OpenVINOManager(BaseBackendManager): +@dataclass +class OpenVINOParam(BaseBackendParam): + """OpenVINO backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + bin_name (str): Serialized bin file. If not given, bin_name would be + the same as file_name with postfix `.param` + input_shapes (ShapeType): The Default shape of the inputs. + output_names (List[str]): Names of the outputs. + mo_options (str): Additional args to OpenVINO Model Optimizer. + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.xml') + bin_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.bin', base_name='file_name') + mo_options: str = '' - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import OpenVINOWrapper - return OpenVINOWrapper( - ir_model_file=backend_files[0], output_names=output_names) + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str) + assert isinstance(self.file_name, str) + param_file_path = osp.join(self.work_dir, self.file_name) + assert isinstance(self.bin_name, str) + bin_file_path = osp.join(self.work_dir, self.bin_name) + return param_file_path, bin_file_path + + +_BackendParam = OpenVINOParam + + +@BACKEND_MANAGERS.register('openvino', param=_BackendParam, ir_param=ONNXParam) +class OpenVINOManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -60,42 +75,188 @@ def get_version(cls) -> str: @classmethod def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + onnx_path: str, + model_path: str, + input_info: Dict[str, Sequence], + output_names: Sequence[str], + bin_path: Optional[str] = None, + work_dir: Optional[str] = None, + mo_options: str = '') -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. + onnx_path (str): The ONNX model files. + model_path (str): The save model path. + input_info (Dict[str, Sequence]): The dictionary about input name + and corresponding shapes. + output_names (str): The output names of the model. + bin_path (str): The save weight path. work_dir (str): The work directory, backend files and logs should be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. - Returns: - Sequence[str]: Backend files. + mo_options (str): Other args and flags that feeds to mo. """ - from . import is_available - assert is_available(), \ + assert cls.is_available(), \ 'OpenVINO is not available, please install OpenVINO first.' + from .onnx2openvino import from_onnx + + if work_dir is None: + with tempfile.TemporaryDirectory() as work_dir: + from_onnx( + onnx_path, + model_path, + input_info=input_info, + output_names=output_names, + bin_path=bin_path, + work_dir=work_dir, + mo_options=mo_options) + else: + from_onnx( + onnx_path, + model_path, + input_info=input_info, + output_names=output_names, + bin_path=bin_path, + work_dir=work_dir, + mo_options=mo_options) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + assert isinstance(param.bin_name, str) + bin_path = osp.join(param.work_dir, param.bin_name) + + input_info = param.input_shapes + output_names = param.output_names + mo_options = param.mo_options + + cls.to_backend( + ir_model, + model_path, + input_info=input_info, + output_names=output_names, + bin_path=bin_path, + work_dir=param.work_dir, + mo_options=mo_options) + + @classmethod + def build_wrapper(cls, + model_path: str, + bin_path: Optional[str] = None, + output_names: Optional[Sequence[str]] = None): + """Build the wrapper for the backend model. + Args: + model_path (str): OpenVINO model path. + bin_path (str): OpenVINO weight path. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + """ + from .wrapper import OpenVINOWrapper + return OpenVINOWrapper( + model_path=model_path, + bin_path=bin_path, + output_names=output_names) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + param_path, bin_path = param.get_model_files() + output_names = param.output_names + if len(output_names) == 0: + output_names = None + return cls.build_wrapper( + param_path, bin_path, output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ from mmdeploy.apis.openvino import (get_input_info_from_cfg, - get_mo_options_from_cfg, - get_output_model_file) + get_mo_options_from_cfg) from mmdeploy.utils import get_ir_config - from .onnx2openvino import from_onnx + ir_config = get_ir_config(config) + output_names = ir_config.get('output_names', []) + input_info = get_input_info_from_cfg(config) + mo_options = get_mo_options_from_cfg(config) + mo_options = mo_options.get_options() + + kwargs.setdefault('work_dir', work_dir) + kwargs.setdefault('input_shapes', input_info) + kwargs.setdefault('output_names', output_names) + kwargs.setdefault('mo_options', mo_options) + + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + if len(backend_files) > 1: + kwargs['bin_name'] = backend_files[1] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + + parsed_args = parser.parse_args(args) + yield parsed_args + + # perform command + command = parsed_args._command + + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + bin_name=parsed_args.bin_name, + input_shapes=parsed_args.input_shapes, + output_names=parsed_args.output_names, + mo_options=parsed_args.mo_options) - openvino_files = [] - for onnx_path in ir_files: - model_xml_path = get_output_model_file(onnx_path, 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) - from_onnx(onnx_path, work_dir, input_info, output_names, - mo_options) - openvino_files.append(model_xml_path) - - return openvino_files + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/openvino/onnx2openvino.py b/mmdeploy/backend/openvino/onnx2openvino.py index 196066ae72..f9299e9043 100644 --- a/mmdeploy/backend/openvino/onnx2openvino.py +++ b/mmdeploy/backend/openvino/onnx2openvino.py @@ -1,5 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. import os.path as osp +import shutil import subprocess import tempfile from subprocess import PIPE, CalledProcessError, run @@ -9,7 +10,6 @@ import onnx from mmdeploy.utils import get_root_logger -from .utils import ModelOptimizerOptions def get_mo_command() -> str: @@ -57,10 +57,12 @@ def get_output_model_file(onnx_path: str, work_dir: str) -> str: def from_onnx(onnx_model: Union[str, onnx.ModelProto], - output_file_prefix: str, + xml_path: str, input_info: Dict[str, Sequence[int]], output_names: Sequence[str], - mo_options: Optional[ModelOptimizerOptions] = None): + bin_path: Optional[str] = None, + work_dir: str = './', + mo_options: str = ''): """Convert ONNX to OpenVINO. Examples: @@ -68,21 +70,25 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], >>> input_info = {'input': [1,3,800,1344]} >>> output_names = ['dets', 'labels'] >>> onnx_path = 'work_dir/end2end.onnx' - >>> output_dir = 'work_dir' - >>> from_onnx( onnx_path, output_dir, input_info, output_names) + >>> xml_path = 'work_dir/end2end.xml' + >>> from_onnx( onnx_path, xml_path, input_info, output_names) Args: onnx_model (str|ModelProto): The onnx model or its path. - output_file_prefix (str): The path to the directory for saving - the results. + xml_path (str): The save model path. input_info (Dict[str, Sequence[int]]): The shape of each input. output_names (Sequence[str]): Output names. Example: ['dets', 'labels']. - mo_options (None | ModelOptimizerOptions): The class with - additional arguments for the Model Optimizer. + bin_path (str): The save weight path. + work_dir (str): The path to the directory for saving + the results. + mo_options (str): Additional arguments for the Model Optimizer. """ - work_dir = output_file_prefix + + if bin_path is None: + bin_path = osp.splitext(xml_path)[0] + '.bin' + input_names = ','.join(input_info.keys()) input_shapes = ','.join(str(list(elem)) for elem in input_info.values()) output = ','.join(output_names) @@ -104,8 +110,7 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], f'--output="{output}" ' \ f'--input="{input_names}" ' \ f'--input_shape="{input_shapes}" ' - if mo_options is not None: - mo_args += mo_options.get_options() + mo_args += mo_options command = f'{mo_command} {mo_args}' @@ -116,4 +121,10 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], logger.debug(mo_output.stderr.decode()) model_xml = get_output_model_file(onnx_path, work_dir) - logger.info(f'Successfully exported OpenVINO model: {model_xml}') + model_bin = osp.splitext(model_xml)[0] + '.bin' + + # move result to save path + shutil.move(model_xml, xml_path) + shutil.move(model_bin, bin_path) + + logger.info(f'Successfully exported OpenVINO model: {xml_path}') diff --git a/mmdeploy/backend/openvino/wrapper.py b/mmdeploy/backend/openvino/wrapper.py index ab91f8331b..a2ca61df9c 100644 --- a/mmdeploy/backend/openvino/wrapper.py +++ b/mmdeploy/backend/openvino/wrapper.py @@ -15,7 +15,8 @@ class OpenVINOWrapper(BaseWrapper): """OpenVINO wrapper for inference in CPU. Args: - ir_model_file (str): Input OpenVINO IR model file. + model_path (str): Input OpenVINO IR model file. + bin_path (str): Input OpenVINO weight file. output_names (Sequence[str] | None): Names of model outputs in order. Defaults to `None` and the wrapper will load the output names from model. @@ -24,22 +25,25 @@ class OpenVINOWrapper(BaseWrapper): >>> from mmdeploy.backend.openvino import OpenVINOWrapper >>> import torch >>> - >>> ir_model_file = 'model.xml' - >>> model = OpenVINOWrapper(ir_model_file) + >>> model_path = 'model.xml' + >>> bin_path = 'bin.xml' + >>> model = OpenVINOWrapper(model_path) >>> inputs = dict(input=torch.randn(1, 3, 224, 224, device='cpu')) >>> outputs = model(inputs) >>> print(outputs) """ def __init__(self, - ir_model_file: str, + model_path: str, + bin_path: Optional[str] = None, output_names: Optional[Sequence[str]] = None, **kwargs): from openvino.inference_engine import IECore self.ie = IECore() - bin_path = osp.splitext(ir_model_file)[0] + '.bin' - self.net = self.ie.read_network(ir_model_file, bin_path) + if bin_path is None: + bin_path = osp.splitext(model_path)[0] + '.bin' + self.net = self.ie.read_network(model_path, bin_path) for input in self.net.input_info.values(): batch_size = input.input_data.shape[0] dims = len(input.input_data.shape) diff --git a/mmdeploy/backend/pplnn/__init__.py b/mmdeploy/backend/pplnn/__init__.py index 8ed6101188..9dbe1f4415 100644 --- a/mmdeploy/backend/pplnn/__init__.py +++ b/mmdeploy/backend/pplnn/__init__.py @@ -1,13 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import PPLNNManager +from .backend_manager import PPLNNManager, PPLNNParam _BackendManager = PPLNNManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['PPLNNManager'] - -if is_available(): - from .utils import register_engines - from .wrapper import PPLNNWrapper - __all__ += ['PPLNNWrapper', 'register_engines'] +__all__ = ['PPLNNParam', 'PPLNNManager'] diff --git a/mmdeploy/backend/pplnn/backend_manager.py b/mmdeploy/backend/pplnn/backend_manager.py index fc331d3042..fec8555d16 100644 --- a/mmdeploy/backend/pplnn/backend_manager.py +++ b/mmdeploy/backend/pplnn/backend_manager.py @@ -1,40 +1,55 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os.path as osp -from typing import Any, Optional, Sequence +from argparse import ArgumentParser +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor, import_custom_modules) -@BACKEND_MANAGERS.register('pplnn') -class PPLNNManager(BaseBackendManager): +@dataclass +class PPLNNParam(BaseBackendParam): + """PPLNN backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + algo_name (str): Serialized algorithm file. If not given, + algo_name would be he same as file_name with postfix `.json` + input_shapes (ShapeType): The Default shape of the inputs. + device (str): Inference device. + disable_avx512 (bool): Whether to disable avx512 for x86. + Defaults to `False`. + quick_select (bool): Whether to use default algorithms. + Defaults to `False`. + """ - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import PPLNNWrapper - return PPLNNWrapper( - onnx_file=backend_files[0], - algo_file=backend_files[1] if len(backend_files) > 1 else None, - device=device, - output_names=output_names) + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.onnx') + algo_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.json', base_name='file_name') + disable_avx512: bool = False + quick_select: bool = False + + def get_model_files(self) -> List[str]: + """get the model files.""" + assert isinstance(self.work_dir, str) + assert isinstance(self.file_name, str) + param_file_path = osp.join(self.work_dir, self.file_name) + assert isinstance(self.algo_name, str) + algorithm_file_path = osp.join(self.work_dir, self.algo_name) + return param_file_path, algorithm_file_path + + +_BackendParam = PPLNNParam + + +@BACKEND_MANAGERS.register('pplnn', param=_BackendParam, ir_param=ONNXParam) +class PPLNNManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -64,45 +79,184 @@ def get_version(cls) -> str: @classmethod def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, + onnx_file: str, + output_file: str, + algo_file: Optional[str] = None, + input_shapes: Optional[Dict[str, Sequence]] = None, device: str = 'cpu', - **kwargs) -> Sequence[str]: + disable_avx512: bool = False, + quick_select: bool = False) -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. + onnx_file (str): Path of input ONNX model file. + output_file (str): Path of output ONNX model file. + algo_file (str): Path of PPLNN algorithm file. + input_shapes (Dict[str, Sequence[int]] | None): Shapes for PPLNN + optimization, default to None. device (str, optional): The device type. Defaults to 'cpu'. + disable_avx512 (bool): Whether to disable avx512 for x86. + Defaults to `False`. + quick_select (bool): Whether to use default algorithms. + Defaults to `False`. Returns: Sequence[str]: Backend files. """ - from mmdeploy.utils import get_model_inputs - from . import is_available from .onnx2pplnn import from_onnx - assert is_available(), \ + assert cls.is_available(), \ 'PPLNN is not available, please install PPLNN first.' - pplnn_files = [] - for onnx_path in ir_files: - algo_file = onnx_path.replace('.onnx', '.json') - model_inputs = get_model_inputs(deploy_cfg) - assert 'opt_shape' in model_inputs, 'Expect opt_shape ' \ - 'in deploy config for PPLNN' - # PPLNN accepts only 1 input shape for optimization, - # may get changed in the future - input_shapes = [model_inputs.opt_shape] - algo_prefix = osp.splitext(algo_file)[0] - from_onnx( - onnx_path, - algo_prefix, - device=device, - input_shapes=input_shapes) - pplnn_files += [onnx_path, algo_file] - - return pplnn_files + from_onnx( + onnx_file, + output_file, + algo_file, + input_shapes=input_shapes, + device=device, + disable_avx512=disable_avx512, + quick_select=quick_select) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + assert isinstance(param.algo_name, str) + algo_path = osp.join(param.work_dir, param.algo_name) + + input_shapes = param.input_shapes + device = param.device + + cls.to_backend( + ir_model, + model_path, + algo_file=algo_path, + input_shapes=input_shapes, + device=device, + disable_avx512=param.disable_avx512, + quick_select=param.quick_select) + + @classmethod + def build_wrapper(cls, + onnx_file: str, + algo_file: Optional[str] = None, + device: str = 'cpu', + output_names: Optional[Sequence[str]] = None): + """Build the wrapper for the backend model. + + Args: + onnx_file (str): Path of input ONNX model file. + algo_file (str): Path of PPLNN algorithm file. + device (str, optional): The device info. Defaults to 'cpu'. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + """ + from .wrapper import PPLNNWrapper + return PPLNNWrapper( + onnx_file=onnx_file, + algo_file=algo_file if osp.exists(algo_file) else None, + device=device, + output_names=output_names) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_path, algo_path = param.get_model_files() + output_names = param.output_names + if len(output_names) == 0: + output_names = None + device = param.device + return cls.build_wrapper( + model_path, algo_path, device=device, output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + from mmdeploy.utils import get_model_inputs + model_inputs = get_model_inputs(config) + input_shapes = model_inputs.get('opt_shape', [1, 3, 224, 224]) + input_shapes = [input_shapes] + + kwargs.setdefault('work_dir', work_dir) + kwargs.setdefault('input_shapes', input_shapes) + + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + if len(backend_files) > 1: + kwargs['algo_name'] = backend_files[1] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Import custom modules.') + + parsed_args = parser.parse_args(args) + yield parsed_args + import_custom_modules(parsed_args.custom_modules) + + # perform command + command = parsed_args._command + + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + algo_name=parsed_args.algo_name, + input_shapes=parsed_args.input_shapes, + device=parsed_args.device, + disable_avx512=parsed_args.disable_avx512, + quick_select=parsed_args.quick_select) + + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/pplnn/onnx2pplnn.py b/mmdeploy/backend/pplnn/onnx2pplnn.py index 1784f1b81b..8e32a3fbc4 100644 --- a/mmdeploy/backend/pplnn/onnx2pplnn.py +++ b/mmdeploy/backend/pplnn/onnx2pplnn.py @@ -1,14 +1,19 @@ # Copyright (c) OpenMMLab. All rights reserved. -from typing import Optional, Sequence +from typing import Dict, Optional, Sequence + +import onnx from mmdeploy.utils.device import parse_cuda_device_id from .utils import create_runtime, register_engines -def from_onnx(onnx_model: str, - output_file_prefix: str, +def from_onnx(onnx_file: str, + output_file: str, + algo_file: Optional[str] = None, device: str = 'cuda:0', - input_shapes: Optional[Sequence[Sequence[int]]] = None, + input_shapes: Optional[Dict[str, Sequence[int]]] = None, + disable_avx512: bool = False, + quick_select: bool = False, **kwargs): """Convert ONNX to PPLNN. PPLNN is capable of optimizing onnx model. The optimized algorithm is saved @@ -17,16 +22,20 @@ def from_onnx(onnx_model: str, our codebase, we only pass one input shape which can be modified by users' own preferences. Args: - output_file_prefix (str): File path to save PPLNN optimization - algorithm and ONNX file - onnx_model (str): Input onnx model. + onnx_file (str): Input onnx model. + output_file (str): Path of output ONNX model file. + algo_file (str): Path of PPLNN algorithm file. device (str): A string specifying device, defaults to 'cuda:0'. - input_shapes (Sequence[Sequence[int]] | None): Shapes for PPLNN + input_shapes (Dict[str, Sequence[int]] | None): Shapes for PPLNN optimization, default to None. + disable_avx512 (bool): Whether to disable avx512 for x86. + Defaults to `False`. + quick_select (bool): Whether to use default algorithms. + Defaults to `False`. Examples: >>> from mmdeploy.apis.pplnn import from_onnx >>> - >>> from_onnx(onnx_model = 'example.onnx', + >>> from_onnx(onnx_file = 'example.onnx', output_file_prefix = 'example') """ if device == 'cpu': @@ -35,19 +44,23 @@ def from_onnx(onnx_model: str, assert 'cuda' in device, f'unexpected device: {device}, must contain ' '`cpu` or `cuda`' device_id = parse_cuda_device_id(device) + + onnx_model = onnx.load(onnx_file) + input_names = [i.name for i in onnx_model.graph.input] + if input_shapes is None: input_shapes = [[1, 3, 224, 224]] # PPLNN default shape for optimization + elif isinstance(input_shapes, Dict): + input_shapes = [input_shapes[name] for name in input_names] - algo_file = output_file_prefix + '.json' - onnx_output_path = output_file_prefix + '.onnx' engines = register_engines( device_id, - disable_avx512=False, - quick_select=False, + disable_avx512=disable_avx512, + quick_select=quick_select, export_algo_file=algo_file, input_shapes=input_shapes) - _ = create_runtime(onnx_model, engines) # side effect: export algorithms + _ = create_runtime(onnx_file, engines) # side effect: export algorithms import shutil - if onnx_output_path != onnx_model: - shutil.copy2(onnx_model, onnx_output_path) + if output_file != onnx_file: + shutil.copy2(onnx_file, output_file) diff --git a/mmdeploy/backend/pplnn/wrapper.py b/mmdeploy/backend/pplnn/wrapper.py index 981a388a69..9a3e73a9c1 100644 --- a/mmdeploy/backend/pplnn/wrapper.py +++ b/mmdeploy/backend/pplnn/wrapper.py @@ -20,7 +20,7 @@ class PPLNNWrapper(BaseWrapper): algo_file (str): Path of PPLNN algorithm file. device_id (int): Device id to put model. Examples: - >>> from mmdeploy.backend.pplnn import PPLNNWrapper + >>> from mmdeploy.backend.pplnn.wrapper import PPLNNWrapper >>> import torch >>> >>> onnx_file = 'model.onnx' diff --git a/mmdeploy/backend/rknn/__init__.py b/mmdeploy/backend/rknn/__init__.py index b0a84a090b..d43bae49a5 100644 --- a/mmdeploy/backend/rknn/__init__.py +++ b/mmdeploy/backend/rknn/__init__.py @@ -1,12 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import RKNNManager +from .backend_manager import RKNNManager, RKNNParam _BackendManager = RKNNManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['RKNNManager'] - -if is_available(): - from .wrapper import RKNNWrapper - __all__ += ['RKNNWrapper'] +__all__ = ['RKNNParam', 'RKNNManager'] diff --git a/mmdeploy/backend/rknn/backend_manager.py b/mmdeploy/backend/rknn/backend_manager.py index 8660aa313c..7a56307df9 100644 --- a/mmdeploy/backend/rknn/backend_manager.py +++ b/mmdeploy/backend/rknn/backend_manager.py @@ -1,43 +1,101 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os.path as osp -from typing import Any, Callable, Optional, Sequence +import re +from argparse import Action, ArgumentParser +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Sequence -from mmdeploy.utils import get_common_config -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor, import_custom_modules) -@BACKEND_MANAGERS.register('rknn') -class RKNNManager(BaseBackendManager): +class MeanStdAction(Action): + """dtype argparse action.""" + + def __call__(self, parser, namespace, values, option_string=None): + """call action.""" + args = values + if isinstance(args, str): + args = [args] + + pattern = r'^[^\S\n]*((?P\w+):)?(?P(.+)(,.+)*)$' + ret = dict() + for arg in args: + if len(arg) == 0: + continue + m = re.match(pattern, arg) + if m is None: + raise ValueError(f'Can not parse value: {arg}') + input_name = m.group('input_name') + val = m.group('val') + val = val.split(',') + val = tuple(float(v) for v in val) + if input_name in ret: + raise NameError(f'value of `{input_name}` has been assigned' + 'more than once.') + ret[input_name] = val + + setattr(namespace, self.dest, ret) + + +@dataclass +class RKNNParam(BaseBackendParam): + """RKNN backend parameters. + + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + input_shapes (ShapeType): The Default shape of the inputs. + input_names (List[str]): Names of the inputs. + output_names (List[str]): Names of the outputs. + mean_values (Dict[str, List[int]]): mean of the inputs. with format: + `input_name1:v0,v1,v2 input_name2:v0,v1,v2` + std_values (Dict[str, List[int]]): mean of the inputs. with format: + `input_name1:v0,v1,v2 input_name2:v0,v1,v2` + device (str): Target platform, such as `rv1126` or `rk3588`. + optimization_level (int): The optimization level of model. Default to 1 + do_quantization (bool): Enable model quantization. + dataset (str): Dataset file. Each line is an image path. + pre_compile (bool): Pre compile the model (smaller size and load + quicker, but can't run on simulator) + """ + + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.rknn') + mean_values: Dict[str, List[int]] = None + std_values: Dict[str, List[int]] = None + optimization_level: int = 1 + do_quantization: bool = False + dataset: str = None + pre_compile: bool = False + + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str), ('Expect string work_dir, ' + f'got {self.work_dir}') + assert isinstance(self.file_name, str), ('Expect string file_name, ' + f'got {self.file_name}') + return osp.join(self.work_dir, self.file_name) @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + def add_argument(cls, parser: ArgumentParser, name: str, dtype: Any, + default: Any, desc: str): + arg_name = f'--{name.replace("_", "-")}' + if name == 'mean_values' or name == 'std_values': + parser.add_argument( + arg_name, action=MeanStdAction, nargs='+', help=desc) + else: + return super().add_argument(parser, name, dtype, default, desc) - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import RKNNWrapper - common_config = get_common_config(deploy_cfg) - return RKNNWrapper( - model=backend_files[0], - common_config=common_config, - input_names=input_names, - output_names=output_names) +_BackendParam = RKNNParam + + +@BACKEND_MANAGERS.register('rknn', param=_BackendParam, ir_param=ONNXParam) +class RKNNManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -116,35 +174,264 @@ def check_env(cls, log_callback: Callable = lambda _: _) -> str: @classmethod def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + onnx_path: str, + output_path: str, + input_names: List[str], + output_names: List[str], + input_shapes: Dict[str, Sequence], + rknn_config: Any, + do_quantization: bool = False, + dataset: Optional[str] = None, + pre_compile: bool = False) -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. + onnx_path (str): The intermediate representation files. + output_path (str): File path to save RKNN model. + input_names (List[str]): Names of the inputs. + output_names (List[str]): Names of the outputs. + input_shapes (ShapeType): The Default shape of the inputs. + rknn_config (RKNNConfig): Config of the rknn toolset. Defined in + `mmdeploy.backend.rknn.onnx2rknn`. + optimization_level (int): The optimization level of model. + Default to 1 + do_quantization (bool): Enable model quantization. + dataset (str): Dataset file. Each line is an image path. + pre_compile (bool): Pre compile the model (smaller size and load + quicker, but can't run on simulator) Returns: Sequence[str]: Backend files. """ - from . import is_available - assert is_available( + assert cls.is_available( ), 'RKNN is not available, please install RKNN first.' from .onnx2rknn import onnx2rknn - backend_files = [] - for model_id, onnx_path in zip(range(len(ir_files)), ir_files): - pre_fix_name = osp.splitext(osp.split(onnx_path)[1])[0] - output_path = osp.join(work_dir, pre_fix_name + '.rknn') - onnx2rknn(onnx_path, output_path, deploy_cfg) - backend_files.append(output_path) + onnx2rknn( + onnx_path, + output_path, + input_names, + output_names, + input_shapes, + rknn_config=rknn_config, + do_quantization=do_quantization, + dataset=dataset, + pre_compile=pre_compile) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + from .onnx2rknn import RKNNConfig + + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + + input_shapes = param.input_shapes + device = param.device + + # get input names + input_names = param.input_names + output_names = param.output_names + mean_values = param.mean_values + if mean_values is not None: + mean_values = list(param.mean_values[name] for name in input_names + if name in param.mean_values) + std_values = param.std_values + if std_values is not None: + std_values = list(param.std_values[name] for name in input_names + if name in param.std_values) + optimization_level = param.optimization_level + target_platform = device + rknn_config = RKNNConfig( + mean_values=mean_values, + std_values=std_values, + optimization_level=optimization_level, + target_platform=target_platform) + + do_quantization = param.do_quantization + dataset = param.dataset + pre_compile = param.pre_compile + + cls.to_backend( + ir_model, + model_path, + input_names, + output_names, + input_shapes=input_shapes, + rknn_config=rknn_config, + do_quantization=do_quantization, + dataset=dataset, + pre_compile=pre_compile) + + @classmethod + def build_wrapper(cls, + model_path: str, + target_platform: str, + input_names: Optional[Sequence[str]] = None, + output_names: Optional[Sequence[str]] = None): + """Build the wrapper for the backend model. + + Args: + model_path (str): Backend model file. + target_platform (str): Target platform, such as `rv1126` or + `rk3588`. + input_names (Optional[Sequence[str]], optional): input names. + Defaults to None. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + """ + from .wrapper import RKNNWrapper + return RKNNWrapper( + model=model_path, + target_platform=target_platform, + input_names=input_names, + output_names=output_names) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_path = param.get_model_files() + input_names = param.input_names + output_names = param.output_names + device = param.device + return cls.build_wrapper( + model_path, + target_platform=device, + input_names=input_names, + output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + from mmdeploy.utils import (get_common_config, get_ir_config, + get_quantization_config) + from mmdeploy.utils.config_utils import get_backend_config + + deploy_cfg = config + common_params = get_common_config(deploy_cfg) + onnx_params = get_ir_config(deploy_cfg) + quantization_cfg = get_quantization_config(deploy_cfg) + + input_names = onnx_params.get('input_names', None) + output_names = onnx_params.get('output_names', None) + input_size_list = get_backend_config(deploy_cfg).get( + 'input_size_list', None) + + mean_values = common_params.get('mean_values', None) + std_values = common_params.get('std_values', None) + target_platform = common_params['target_platform'] + optimization_level = common_params['optimization_level'] + if mean_values is not None: + mean_values = dict(zip(input_names, mean_values)) + if std_values is not None: + std_values = dict(zip(input_names, std_values)) + + do_quantization = quantization_cfg.get('do_quantization', False) + dataset = quantization_cfg.get('dataset', None) + pre_compile = quantization_cfg.get('pre_compile', False) + rknn_batch_size = quantization_cfg.get('rknn_batch_size', -1) + + batched_input_size_list = list( + (rknn_batch_size, *size) for size in input_size_list) + input_shapes = dict(zip(input_names, batched_input_size_list)) + + kwargs.setdefault('input_names', input_names) + kwargs.setdefault('output_names', output_names) + kwargs.setdefault('input_shapes', input_shapes) + + kwargs.setdefault('mean_values', mean_values) + kwargs.setdefault('std_values', std_values) + kwargs['device'] = target_platform + kwargs.setdefault('optimization_level', optimization_level) + + kwargs.setdefault('do_quantization', do_quantization) + kwargs.setdefault('dataset', dataset) + kwargs.setdefault('pre_compile', pre_compile) + + kwargs.setdefault('work_dir', work_dir) + kwargs.setdefault('input_shapes', input_shapes) + + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Import custom modules.') + + parsed_args = parser.parse_args(args) + yield parsed_args + import_custom_modules(parsed_args.custom_modules) + + # perform command + command = parsed_args._command + + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + input_shapes=parsed_args.input_shapes, + input_names=parsed_args.input_names, + output_names=parsed_args.output_names, + mean_values=parsed_args.mean_values, + std_values=parsed_args.std_values, + device=parsed_args.device, + optimization_level=parsed_args.optimization_level, + do_quantization=parsed_args.do_quantization, + dataset=parsed_args.dataset, + pre_compile=parsed_args.pre_compile) - return backend_files + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/rknn/onnx2rknn.py b/mmdeploy/backend/rknn/onnx2rknn.py index a0d0583567..c809b99ea2 100644 --- a/mmdeploy/backend/rknn/onnx2rknn.py +++ b/mmdeploy/backend/rknn/onnx2rknn.py @@ -1,14 +1,18 @@ # Copyright (c) OpenMMLab. All rights reserved. -from typing import Optional, Union +import os.path as osp +from dataclasses import asdict, dataclass +from typing import Dict, List, Optional -import mmengine -from rknn.api import RKNN +from mmdeploy.utils import get_root_logger -from mmdeploy.utils import (get_common_config, get_normalization, - get_onnx_config, get_partition_config, - get_quantization_config, get_rknn_quantization, - get_root_logger, load_config) -from mmdeploy.utils.config_utils import get_backend_config + +@dataclass +class RKNNConfig: + """RKNN Config.""" + mean_values: List[List[int]] = None + std_values: List[List[int]] = None + optimization_level: int = 1 + target_platform: str = None def rknn_package_info(): @@ -24,10 +28,13 @@ def rknn_package_info(): def onnx2rknn(onnx_model: str, output_path: str, - deploy_cfg: Union[str, mmengine.Config], - model_cfg: Optional[Union[str, mmengine.Config]] = None, - dataset_file: Optional[str] = None, - **kwargs): + input_names: List[str], + output_names: List[str], + input_shapes: Dict[str, List], + rknn_config: RKNNConfig, + do_quantization: bool = False, + dataset: Optional[str] = None, + pre_compile: bool = False): """Convert ONNX to RKNN. RKNN-Toolkit2 is a software development kit for users to perform model @@ -37,39 +44,27 @@ def onnx2rknn(onnx_model: str, Args: onnx_model (str): Input onnx model. output_path (str): File path to save RKNN model. - deploy_cfg (str | mmengine.Config): The path or content of config. - model_cfg (str | mmengine.Config): The path or content of model config. - dataset_file (str | None): The dataset file for quatization. Default to - None. + input_names (List[str]): Names of the inputs. + output_names (List[str]): Names of the outputs. + input_shapes (ShapeType): The Default shape of the inputs. + rknn_config (RKNNConfig): Config of the rknn toolset. Defined in + `mmdeploy.backend.rknn.onnx2rknn`. + do_quantization (bool): Enable model quantization. + dataset (str): Dataset file. Each line is an image path. + pre_compile (bool): Pre compile the model (smaller size and load + quicker, but can't run on simulator) """ + from rknn.api import RKNN logger = get_root_logger() - # load deploy_cfg if necessary - deploy_cfg = load_config(deploy_cfg)[0] - - common_params = get_common_config(deploy_cfg) - onnx_params = get_onnx_config(deploy_cfg) - quantization_cfg = get_quantization_config(deploy_cfg) - - input_names = onnx_params.get('input_names', None) - output_names = onnx_params.get('output_names', None) - input_size_list = get_backend_config(deploy_cfg).get( - 'input_size_list', None) - # update norm value - if get_rknn_quantization(deploy_cfg) is True and model_cfg is not None: - transform = get_normalization(model_cfg) - common_params.update( - dict( - mean_values=[transform['mean']], - std_values=[transform['std']])) - - # update output_names for partition models - if get_partition_config(deploy_cfg) is not None: - import onnx - _onnx_model = onnx.load(onnx_model) - output_names = [node.name for node in _onnx_model.graph.output] + # get input/output names + input_size_list = [list(input_shapes[name][1:]) for name in input_names] + + # init rknn rknn = RKNN(verbose=True) - rknn.config(**common_params) + rknn.config(**asdict(rknn_config)) + + # load onnx ret = rknn.load_onnx( model=onnx_model, inputs=input_names, @@ -79,19 +74,34 @@ def onnx2rknn(onnx_model: str, logger.error('Load model failed!') exit(ret) - dataset_cfg = quantization_cfg.get('dataset', None) - if dataset_cfg is None: - quantization_cfg.update(dict(dataset=dataset_file)) - if dataset_file is None: - quantization_cfg.update(dict(do_quantization=False)) + # quantization + quantization_cfg = dict() + if do_quantization: + # disable quantization if dataset not exist + if not osp.exists(dataset): + do_quantization = False logger.warning('no dataset passed in, quantization is skipped') - if rknn_package_info()['name'] == 'rknn-toolkit2': - quantization_cfg.pop('pre_compile', None) + else: + quantization_cfg['dataset'] = dataset + + # set batch size + if do_quantization: + batch_size = input_size_list[0][0] + assert all(batch_size == shape[0] for shape in input_size_list) + quantization_cfg['rknn_batch_size'] = batch_size + + # set pre compile + if rknn_package_info()['name'] != 'rknn-toolkit2': + quantization_cfg['pre_compile'] = pre_compile + + # do quantization + quantization_cfg['do_quantization'] = do_quantization ret = rknn.build(**quantization_cfg) if ret != 0: logger.error('Build model failed!') exit(ret) + # export ret = rknn.export_rknn(output_path) if ret != 0: logger.error('Export rknn model failed!') diff --git a/mmdeploy/backend/rknn/wrapper.py b/mmdeploy/backend/rknn/wrapper.py index 70de61a2d9..8ec172e273 100644 --- a/mmdeploy/backend/rknn/wrapper.py +++ b/mmdeploy/backend/rknn/wrapper.py @@ -22,7 +22,7 @@ class RKNNWrapper(BaseWrapper): verbose (bool): Whether verbose during inference. Examples: - >>> from mmdeploy.backend.rknn import RKNNWrapper + >>> from mmdeploy.backend.rknn.wrapper import RKNNWrapper >>> import torch >>> >>> model = 'model.rknn' @@ -34,17 +34,16 @@ class RKNNWrapper(BaseWrapper): def __init__(self, model: str, - common_config: Dict = dict(target_platform=None), + target_platform: str, input_names: Optional[Sequence[str]] = None, output_names: Optional[Sequence[str]] = None, - verbose=True, - **kwargs): + verbose=True): logger = get_root_logger() # Create RKNN object self.rknn = RKNN(verbose=verbose) self.rknn.load_rknn(model) self.input_names = input_names - ret = self.rknn.init_runtime(target=common_config['target_platform']) + ret = self.rknn.init_runtime(target=target_platform) if ret != 0: logger.error('Init runtime environment failed!') exit(ret) diff --git a/mmdeploy/backend/sdk/__init__.py b/mmdeploy/backend/sdk/__init__.py index 46c750c40c..36742db67a 100644 --- a/mmdeploy/backend/sdk/__init__.py +++ b/mmdeploy/backend/sdk/__init__.py @@ -1,16 +1,9 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import SDKManager +from .backend_manager import SDKManager, SDKParam _BackendManager = SDKManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param -__all__ = ['SDKManager'] - -if is_available(): - - try: - from .wrapper import SDKWrapper - __all__ += ['SDKWrapper'] - except Exception: - pass +__all__ = ['SDKParam', 'SDKManager'] diff --git a/mmdeploy/backend/sdk/backend_manager.py b/mmdeploy/backend/sdk/backend_manager.py index e377628234..303bd64123 100644 --- a/mmdeploy/backend/sdk/backend_manager.py +++ b/mmdeploy/backend/sdk/backend_manager.py @@ -2,10 +2,11 @@ import importlib import os.path as osp import sys -from typing import Any, Optional, Sequence +from dataclasses import dataclass +from typing import Any, List from mmdeploy.utils import get_file_path -from ..base import BACKEND_MANAGERS, BaseBackendManager +from ..base import BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam _is_available = False @@ -26,36 +27,50 @@ _is_available = True -@BACKEND_MANAGERS.register('sdk') +@dataclass +class SDKParam(BaseBackendParam): + """SDK backend parameters. + + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + task_name (str): The name of the SDK task. + device (str): Inference device. + """ + task_name: str = None + + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str) + assert isinstance(self.file_name, str) + model_path = osp.join(self.work_dir, self.file_name) + return model_path + + +_BackendParam = SDKParam + + +@BACKEND_MANAGERS.register('sdk', param=SDKParam) class SDKManager(BaseBackendManager): @classmethod def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): + backend_model: str, + task_name: str, + device: str = 'cpu'): """Build the wrapper for the backend model. Args: - backend_files (Sequence[str]): Backend files. + backend_model (str): Backend model. + task_name (str): The name of the SDK task. device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. deploy_cfg (Optional[Any], optional): The deploy config. Defaults to None. """ - assert deploy_cfg is not None, \ - 'Building SDKWrapper requires deploy_cfg' - from mmdeploy.backend.sdk import SDKWrapper - from mmdeploy.utils import SDK_TASK_MAP, get_task_type - task_name = SDK_TASK_MAP[get_task_type(deploy_cfg)]['cls_name'] + from .wrapper import SDKWrapper return SDKWrapper( - model_file=backend_files[0], task_name=task_name, device=device) + model_file=backend_model, task_name=task_name, device=device) @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -81,3 +96,29 @@ def get_version(cls) -> str: return pkg_resources.get_distribution('mmdeploy').version except Exception: return 'None' + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_path = param.get_model_files() + device = param.device + task_name = param.task_name + return cls.build_wrapper( + model_path, task_name=task_name, device=device) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: List[str] = None, + **kwargs) -> _BackendParam: + from mmdeploy.utils import SDK_TASK_MAP, get_task_type + task_name = SDK_TASK_MAP[get_task_type(config)]['cls_name'] + + kwargs.setdefault('task_name', task_name) + return _BackendParam( + work_dir=work_dir, file_name=backend_files[0], **kwargs) diff --git a/mmdeploy/backend/snpe/__init__.py b/mmdeploy/backend/snpe/__init__.py index 5398d0e94a..0608f070a3 100644 --- a/mmdeploy/backend/snpe/__init__.py +++ b/mmdeploy/backend/snpe/__init__.py @@ -1,17 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import SNPEManager -from .onnx2dlc import from_onnx +from .backend_manager import SNPEManager, SNPEParam _BackendManager = SNPEManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['from_onnx', 'SNPEManager'] -if is_available(): - try: - from .wrapper import SNPEWrapper - - __all__ += ['SNPEWrapper'] - except Exception as e: - print(e) - pass +__all__ = ['SNPEParam', 'SNPEManager'] diff --git a/mmdeploy/backend/snpe/backend_manager.py b/mmdeploy/backend/snpe/backend_manager.py index b0ddf2dee7..dca7778790 100644 --- a/mmdeploy/backend/snpe/backend_manager.py +++ b/mmdeploy/backend/snpe/backend_manager.py @@ -1,43 +1,47 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os import os.path as osp import sys -from typing import Any, Optional, Sequence +from argparse import ArgumentParser +from dataclasses import dataclass +from subprocess import call, run +from typing import Any, List, Optional, Sequence +from mmdeploy.ir.onnx import ONNXParam from mmdeploy.utils import get_root_logger -from ..base import BACKEND_MANAGERS, BaseBackendManager +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor, import_custom_modules) -@BACKEND_MANAGERS.register('snpe') -class SNPEManager(BaseBackendManager): +@dataclass +class SNPEParam(BaseBackendParam): + """SNPE backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + uri (str): The uri of remote device. + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.dlc') - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import SNPEWrapper - uri = None - if 'uri' in kwargs: - uri = kwargs['uri'] - return SNPEWrapper( - dlc_file=backend_files[0], uri=uri, output_names=output_names) + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str), ('Expect string work_dir, ' + f'got {self.work_dir}') + assert isinstance(self.file_name, str), ('Expect string file_name, ' + f'got {self.file_name}') + file_name = self.file_name + return osp.join(self.work_dir, file_name) + + +_BackendParam = SNPEParam + + +@BACKEND_MANAGERS.register('snpe', param=_BackendParam, ir_param=ONNXParam) +class SNPEManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -52,48 +56,161 @@ def is_available(cls, with_custom_ops: bool = False) -> bool: onnx2dlc = get_onnx2dlc_path() if onnx2dlc is None: return False - return osp.exists(onnx2dlc) + if not osp.exists(onnx2dlc): + return False + + ret_code = call([onnx2dlc, '-v'], + stdout=open(os.devnull, 'wb'), + stderr=open(os.devnull, 'wb')) + return ret_code == 0 @classmethod - def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - log_level: int = logging.INFO, - device: str = 'cpu', - uri: str = '', - **kwargs) -> Sequence[str]: + def get_version(cls) -> str: + """Get the version of the backend.""" + from .onnx2dlc import get_onnx2dlc_path + onnx2dlc = get_onnx2dlc_path() + snpe_net_run_path = osp.join(osp.split(onnx2dlc)[0], 'snpe-net-run') + if not osp.exists(snpe_net_run_path): + return '' + + command = [snpe_net_run_path, '--version'] + result = run( + command, + stdout=open(os.devnull, 'wb'), + stderr=open(os.devnull, 'wb'), + universal_newlines=True) + if result.returncode != 0: + return '' + else: + return result.stdout[5:] + + @classmethod + def to_backend(cls, onnx_path: str, save_path: str) -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. - Returns: - Sequence[str]: Backend files. + onnx_path (str): The ONNX model to be converted. + output_path (str): The output file. """ - from . import is_available + from .onnx2dlc import from_onnx logger = get_root_logger() - if not is_available(): + if not cls.is_available(): logger.error('snpe support is not available, please check\n' '1) `snpe-onnx-to-dlc` existed in `PATH`\n' '2) snpe only support\n' 'ubuntu18.04') sys.exit(1) + from_onnx(onnx_path, save_path) - from mmdeploy.apis.snpe import get_env_key, get_output_model_file - from .onnx2dlc import from_onnx + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = param.get_model_files() + + cls.to_backend(ir_model, model_path) + + @classmethod + def build_wrapper(cls, model_path: str, uri: Optional[str] = None): + """Build the wrapper for the backend model. - if get_env_key() not in os.environ: + Args: + model_path (str): Backend files. + uri (str): device uri. + """ + from .wrapper import SNPEWrapper + return SNPEWrapper(dlc_file=model_path, uri=uri) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_path = param.get_model_files() + return cls.build_wrapper(model_path, uri=param.uri) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + + uri = kwargs.get('uri', None) + + from .onnx2dlc import get_env_key + + if uri is not None and get_env_key() not in os.environ: os.environ[get_env_key()] = uri - backend_files = [] - for onnx_path in ir_files: - dlc_path = get_output_model_file(onnx_path, work_dir) - onnx_name = osp.splitext(osp.split(onnx_path)[1])[0] - from_onnx(onnx_path, osp.join(work_dir, onnx_name)) - backend_files += [dlc_path] + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Import custom modules.') + + parsed_args = parser.parse_args(args) + yield parsed_args + import_custom_modules(parsed_args.custom_modules) + + # perform command + command = parsed_args._command + + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + uri=parsed_args.uri) - return backend_files + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/snpe/onnx2dlc.py b/mmdeploy/backend/snpe/onnx2dlc.py index 45e727e459..bac6e31dd5 100644 --- a/mmdeploy/backend/snpe/onnx2dlc.py +++ b/mmdeploy/backend/snpe/onnx2dlc.py @@ -47,8 +47,7 @@ def get_output_model_file(onnx_path: str, return save_dlc -def from_onnx(onnx_model: Union[onnx.ModelProto, str], - output_file_prefix: str): +def from_onnx(onnx_model: Union[onnx.ModelProto, str], save_path: str): """Convert ONNX to dlc. We need to use a executable program to convert the `.onnx` file to a `.dlc` @@ -56,12 +55,12 @@ def from_onnx(onnx_model: Union[onnx.ModelProto, str], Example: >>> from mmdeploy.apis.snpe import from_onnx >>> onnx_path = 'work_dir/end2end.onnx' - >>> output_file_prefix = 'work_dir/end2end' - >>> from_onnx(onnx_path, output_file_prefix) + >>> save_path = 'work_dir/end2end' + >>> from_onnx(onnx_path, save_path) Args: onnx_path (ModelProto|str): The path of the onnx model. - output_file_prefix (str): The path to save the output .dlc file. + save_path (str): The path to save the output .dlc file. """ if not isinstance(onnx_model, str): @@ -70,7 +69,7 @@ def from_onnx(onnx_model: Union[onnx.ModelProto, str], else: onnx_path = onnx_model - save_dlc = output_file_prefix + '.dlc' + save_dlc = save_path onnx2dlc = get_onnx2dlc_path() ret_code = call( diff --git a/mmdeploy/backend/snpe/wrapper.py b/mmdeploy/backend/snpe/wrapper.py index f16d6a554b..2a0bac92c6 100644 --- a/mmdeploy/backend/snpe/wrapper.py +++ b/mmdeploy/backend/snpe/wrapper.py @@ -3,7 +3,7 @@ import os import time from random import randint -from typing import Dict, Optional, Sequence, Tuple +from typing import Dict, Optional, Tuple import grpc import inference_pb2 @@ -98,12 +98,10 @@ class SNPEWrapper(BaseWrapper): Args: dlc_file (str): Path of a weight file. - output_names (Sequence[str] | None): Names of model outputs in order. - Defaults to `None` and the wrapper will load the output names from - snpe model. + uri (str): URI of the device Examples: - >>> from mmdeploy.backend.snpe import SNPEWrapper + >>> from mmdeploy.backend.snpe.wrapper import SNPEWrapper >>> import torch >>> >>> snple_file = 'alexnet.dlc' @@ -113,11 +111,7 @@ class SNPEWrapper(BaseWrapper): >>> print(outputs) """ - def __init__(self, - dlc_file: str, - uri: str, - output_names: Optional[Sequence[str]] = None, - **kwargs): + def __init__(self, dlc_file: str, uri: str): logger = get_root_logger() diff --git a/mmdeploy/backend/tensorrt/__init__.py b/mmdeploy/backend/tensorrt/__init__.py index 6587a3ff8c..fd1111e046 100644 --- a/mmdeploy/backend/tensorrt/__init__.py +++ b/mmdeploy/backend/tensorrt/__init__.py @@ -1,24 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. # flake8: noqa -from .backend_manager import TensorRTManager +from .backend_manager import TensorRTManager, TensorRTParam from .init_plugins import load_tensorrt_plugin _BackendManager = TensorRTManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['TensorRTManager'] - -if is_available(): - from .utils import from_onnx, load, save - - __all__ += ['from_onnx', 'save', 'load', 'load_tensorrt_plugin'] - - try: - # import wrapper if pytorch is available - from .torch_allocator import TorchAllocator - from .wrapper import TRTWrapper - __all__ += ['TRTWrapper'] - __all__ += ['TorchAllocator', 'TRTWrapper'] - except Exception: - pass +__all__ = ['load_tensorrt_plugin', 'TensorRTManager', 'TensorRTParam'] diff --git a/mmdeploy/backend/tensorrt/backend_manager.py b/mmdeploy/backend/tensorrt/backend_manager.py index 912d9cf4f0..8b36d26890 100644 --- a/mmdeploy/backend/tensorrt/backend_manager.py +++ b/mmdeploy/backend/tensorrt/backend_manager.py @@ -1,37 +1,73 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os.path as osp -from typing import Any, Callable, Optional, Sequence +import re +from argparse import ArgumentParser +from collections import OrderedDict +from dataclasses import dataclass +from typing import (Any, Callable, Dict, Iterable, List, Optional, Sequence, + Union) -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor, import_custom_modules) -@BACKEND_MANAGERS.register('tensorrt') -class TensorRTManager(BaseBackendManager): +@dataclass +class TensorRTParam(BaseBackendParam): + """TensorRT backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + input_shapes (ShapeType): The Default shape of the inputs. + min_shapes (ShapeType): The minimal shape of the inputs. + max_shapes (ShapeType): The maximal shape of the inputs. + device (str): Device used to perform inference. + fp16_mode (bool): Enable fp16 mode. + int8_mode (bool): Enable int8 quantization. Can be co-exist with + fp16 mode. + int8_algorithm (str): The quantization algorithm, choice from + [`entropy`, `maxmin`] + quanti_data (Union[Iterable, str]): Iterable object to provide the + quantization data. Each iteration gives a dict of input name and + correspond tensor. + max_workspace_size (int): Extra workspace size required by the model. + default to 1Gb. + """ - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.engine') + device: str = 'cuda' + fp16_mode: bool = False + int8_mode: bool = False + int8_algorithm: str = 'entropy' + max_workspace_size: int = 1 << 30 - from .wrapper import TRTWrapper - return TRTWrapper(engine=backend_files[0], output_names=output_names) + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str), ('Expect string work_dir, ' + f'got {self.work_dir}') + assert isinstance(self.file_name, str), ('Expect string file_name, ' + f'got {self.file_name}') + return osp.join(self.work_dir, self.file_name) + + def check_param(self): + """check param validation.""" + super().check_param() + + if self.int8_mode: + if self.int8_algorithm.lower() not in ['entropy', 'minmax']: + raise ValueError( + f'Unsupported int8 algorithm: {self.int8_algorithm}') + + +_BackendParam = TensorRTParam + + +@BACKEND_MANAGERS.register('tensorrt', param=_BackendParam, ir_param=ONNXParam) +class TensorRTManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -85,54 +121,274 @@ def check_env(cls, log_callback: Callable = lambda _: _) -> str: @classmethod def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + ir_path: str, + save_path: str, + input_shapes: Dict[str, Sequence], + min_shapes: Optional[Dict[str, Sequence]] = None, + max_shapes: Optional[Dict[str, Sequence]] = None, + max_workspace_size: int = 0, + fp16_mode: bool = False, + int8_mode: bool = False, + int8_algorithm: str = 'entropy', + calib_data: Optional[Union[str, Iterable]] = None, + device_id: int = 0, + log_level: Any = None): """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. + ir_path (str or onnx.ModelProto): Input ir model to convert from. + save_path (str): The path to save the output model. + input_shapes (Dict[str, Sequence]): The input shapes of + each input. + min_shapes (Dict[str, Sequence]): The min shapes of each input. + max_shapes (Dict[str, Sequence]): The max shapes of each input. + max_workspace_size (int): To set max workspace size of TensorRT + engine. some tactics and layers need large workspace. + fp16_mode (bool): Specifying whether to enable fp16 mode. + Defaults to `False`. + int8_mode (bool): Specifying whether to enable int8 mode. + Defaults to `False`. + int8_algorithm (str): algorithm used to perform the calibration. + calib_data (Iterable|str): An iterable object to provide the input + data. Or qual name of the object. + device_id (int): Choice the device to create engine + log_level (trt.Logger.Severity): The log level of TensorRT. + """ + import tensorrt as trt + + from .utils import from_onnx + if log_level is None: + log_level = trt.Logger.ERROR + + # fill shapes + if min_shapes is None: + min_shapes = input_shapes + if max_shapes is None: + max_shapes = input_shapes + + merged_shapes = OrderedDict() + for name, val in input_shapes.items(): + if name not in min_shapes: + min_shapes[name] = val + if name not in max_shapes: + max_shapes[name] = val + + merged_shapes[name] = dict( + opt_shape=val, + min_shape=min_shapes[name], + max_shape=max_shapes[name]) + + int8_param = dict() + if int8_mode: + if int8_algorithm.lower() == 'entropy': + int8_algo = trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2 + elif int8_algorithm.lower() == 'minmax': + int8_algo = trt.CalibrationAlgoType.MINMAX_CALIBRATION + else: + raise ValueError( + f'Unsupported int8 algorithm: {int8_algorithm}') + + if isinstance(calib_data, str): + from ..base import get_obj_by_qualname + calib_data = get_obj_by_qualname(calib_data) + + int8_param = dict(calib_file=calib_data, algorithm=int8_algo) + + # export model + from_onnx( + ir_path, + save_path, + input_shapes=merged_shapes, + max_workspace_size=max_workspace_size, + fp16_mode=fp16_mode, + int8_mode=int8_mode, + int8_param=int8_param, + device_id=device_id, + log_level=log_level) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + param.check_param() + + assert isinstance(param, _BackendParam), ('Expect _BackendParam ' + f'get {type(param)}') + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + save_path = osp.join(param.work_dir, param.file_name) + input_shapes = param.input_shapes + min_shapes = param.min_shapes + max_shapes = param.max_shapes + max_workspace_size = param.max_workspace_size + fp16_mode = param.fp16_mode + int8_mode = param.int8_mode + device = param.device + + m = re.match(r'^(cuda|CUDA)(:(?P[0-9]+))?$', device) + assert m is not None, f'Unsupported device {device}' + device_id = m.groupdict().get('device_id', 0) + + cls.to_backend( + ir_model, + save_path, + input_shapes=input_shapes, + min_shapes=min_shapes, + max_shapes=max_shapes, + max_workspace_size=max_workspace_size, + fp16_mode=fp16_mode, + int8_mode=int8_mode, + int8_algorithm=param.int8_algorithm, + calib_data=param.quanti_data, + device_id=device_id) + + @classmethod + def build_wrapper( + cls, + engine_path: str, + output_names: Optional[Sequence[str]] = None, + ): + """Build the wrapper for the backend model. + + Args: + engine_path (str): TensorRT engine file. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + """ + + from .wrapper import TRTWrapper + return TRTWrapper(engine=engine_path, output_names=output_names) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + output_names = param.output_names + if output_names is not None and len(output_names) == 0: + output_names = None + return cls.build_wrapper(model_path, output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: List[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + Returns: - Sequence[str]: Backend files. + BaseBackendParam: The packed backend parameter. """ - import os.path as osp - - from mmdeploy.utils import get_model_inputs, get_partition_config - model_params = get_model_inputs(deploy_cfg) - partition_cfgs = get_partition_config(deploy_cfg) - assert len(model_params) == len(ir_files) - - from . import is_available - assert is_available(), ( - 'TensorRT is not available,' - ' please install TensorRT and build TensorRT custom ops first.') - - from .onnx2tensorrt import onnx2tensorrt - backend_files = [] - for model_id, model_param, onnx_path in zip( - range(len(ir_files)), model_params, ir_files): - onnx_name = osp.splitext(osp.split(onnx_path)[1])[0] - save_file = model_param.get('save_file', onnx_name + '.engine') - - partition_type = 'end2end' if partition_cfgs is None \ - else onnx_name - onnx2tensorrt( - work_dir, - save_file, - model_id, - deploy_cfg, - onnx_path, - device=device, - partition_type=partition_type) - - backend_files.append(osp.join(work_dir, save_file)) - - return backend_files + from mmdeploy.utils import (get_calib_config, get_common_config, + get_model_inputs) + common_config = get_common_config(config) + model_inputs = get_model_inputs(config) + + # get shapes + assert len(model_inputs) == 1, ('Can not create param with ' + 'len(model_inputs) > 1') + shapes = model_inputs[0].get('input_shapes', {}) + min_shapes = OrderedDict() + max_shapes = OrderedDict() + input_shapes = OrderedDict() + for name, vals in shapes.items(): + min_shapes[name] = vals.get('min_shape', []) + input_shapes[name] = vals.get('opt_shape', []) + max_shapes[name] = vals.get('max_shape', []) + + # others + max_workspace_size = common_config.get('max_workspace_size', 0) + fp16_mode = common_config.get('fp16_mode', False) + int8_mode = common_config.get('int8_mode', False) + + kwargs.setdefault('min_shapes', min_shapes) + kwargs.setdefault('max_shapes', max_shapes) + kwargs.setdefault('input_shapes', input_shapes) + kwargs.setdefault('max_workspace_size', max_workspace_size) + kwargs.setdefault('fp16_mode', fp16_mode) + kwargs.setdefault('int8_mode', int8_mode) + + if int8_mode: + calib_config = get_calib_config(config) + if calib_config is not None and calib_config.get( + 'create_calib', False): + from ..base import create_h5pydata_generator + assert 'calib_file' in calib_config + calib_path = osp.join(work_dir, calib_config['calib_file']) + calib_data = create_h5pydata_generator(calib_path, + input_shapes) + kwargs.setdefault('quanti_data', calib_data) + + ret = _BackendParam( + work_dir=work_dir, file_name=backend_files[0], **kwargs) + return ret + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Import custom modules.') + + parsed_args = parser.parse_args(args) + yield parsed_args + + # perform command + command = parsed_args._command + + if command == 'convert': + # convert model + import_custom_modules(parsed_args.custom_modules) + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + device=parsed_args.device, + min_shapes=parsed_args.min_shapes, + input_shapes=parsed_args.input_shapes, + max_shapes=parsed_args.max_shapes, + max_workspace_size=parsed_args.max_workspace_size, + fp16_mode=parsed_args.fp16_mode, + int8_mode=parsed_args.int8_mode, + int8_algorithm=parsed_args.int8_algorithm, + quanti_data=parsed_args.quanti_data) + + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/tensorrt/calib_utils.py b/mmdeploy/backend/tensorrt/calib_utils.py index 1c89366fa4..2b402dade9 100644 --- a/mmdeploy/backend/tensorrt/calib_utils.py +++ b/mmdeploy/backend/tensorrt/calib_utils.py @@ -1,5 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -from typing import Any, Dict, Sequence, Union +from typing import Iterable, Sequence import numpy as np import pycuda.autoinit # noqa:F401 @@ -9,7 +9,7 @@ DEFAULT_CALIBRATION_ALGORITHM = trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2 -class HDF5Calibrator(trt.IInt8Calibrator): +class IteratorCalibrator(trt.IInt8Calibrator): """HDF5 calibrator. Args: @@ -24,73 +24,62 @@ class HDF5Calibrator(trt.IInt8Calibrator): def __init__( self, - calib_file: Union[str, Any], - input_shapes: Dict[str, Sequence[int]], - model_type: str = 'end2end', + data_iter: Iterable, device_id: int = 0, algorithm: trt.CalibrationAlgoType = DEFAULT_CALIBRATION_ALGORITHM, **kwargs): super().__init__() - import h5py + self._data_generator = data_iter - if isinstance(calib_file, str): - calib_file = h5py.File(calib_file, mode='r') - - assert 'calib_data' in calib_file - calib_data = calib_file['calib_data'] - assert model_type in calib_data - calib_data = calib_data[model_type] - - self.calib_file = calib_file - self.calib_data = calib_data - self.device_id = device_id - self.algorithm = algorithm - self.input_shapes = input_shapes - self.kwargs = kwargs + self._device_id = device_id + self._algorithm = algorithm # create buffers that will hold data batches - self.buffers = dict() + self._buffers = dict() - self.count = 0 - first_input_group = calib_data[list(calib_data.keys())[0]] - self.dataset_length = len(first_input_group) - self.batch_size = first_input_group['0'].shape[0] + next_data = next(self._data_generator) + names = list(next_data.keys()) + self._batch_size = next_data[names[0]].shape[0] + self._next_data = next_data def __del__(self): """Close h5py file if necessary.""" - if hasattr(self, 'calib_file'): - self.calib_file.close() + del self._data_generator def get_batch(self, names: Sequence[str], **kwargs) -> list: """Get batch data.""" - if self.count < self.dataset_length: + if self._next_data is not None: + # host to device ret = [] - for name in names: - input_group = self.calib_data[name] - data_np = input_group[str(self.count)][...].astype(np.float32) - - # tile the tensor so we can keep the same distribute - opt_shape = self.input_shapes[name]['opt_shape'] - data_shape = data_np.shape - reps = [ - int(np.ceil(opt_s / data_s)) - for opt_s, data_s in zip(opt_shape, data_shape) - ] - - data_np = np.tile(data_np, reps) - - slice_list = tuple(slice(0, end) for end in opt_shape) - data_np = data_np[slice_list] - - data_np_cuda_ptr = cuda.mem_alloc(data_np.nbytes) - cuda.memcpy_htod(data_np_cuda_ptr, - np.ascontiguousarray(data_np)) - self.buffers[name] = data_np_cuda_ptr + for name in names: + data_np = self._next_data[name] + + is_torch_data = False + try: + import torch + if isinstance(data_np, torch.Tensor): + is_torch_data = True + except Exception: + pass + + if is_torch_data: + data_np = data_np.cuda(self._device_id) + self._buffers[name] = data_np + ret.append(data_np.data_ptr()) + else: + assert isinstance(data_np, np.ndarray) + data_np_cuda_ptr = cuda.mem_alloc(data_np.nbytes) + cuda.memcpy_htod(data_np_cuda_ptr, + np.ascontiguousarray(data_np)) + self._buffers[name] = data_np_cuda_ptr + ret.append(data_np_cuda_ptr) + try: + self._next_data = next(self._data_generator) + except StopIteration: + self._next_data = None - ret.append(self.buffers[name]) - self.count += 1 return ret else: return None @@ -101,7 +90,7 @@ def get_algorithm(self) -> trt.CalibrationAlgoType: Returns: trt.CalibrationAlgoType: Calibration algo type. """ - return self.algorithm + return self._algorithm def get_batch_size(self) -> int: """Get batch size. @@ -109,7 +98,7 @@ def get_batch_size(self) -> int: Returns: int: An integer represents batch size. """ - return self.batch_size + return self._batch_size def read_calibration_cache(self, *args, **kwargs): """Read calibration cache. diff --git a/mmdeploy/backend/tensorrt/onnx2tensorrt.py b/mmdeploy/backend/tensorrt/onnx2tensorrt.py index caeddf8e57..fa5b69691e 100644 --- a/mmdeploy/backend/tensorrt/onnx2tensorrt.py +++ b/mmdeploy/backend/tensorrt/onnx2tensorrt.py @@ -78,7 +78,7 @@ def onnx2tensorrt(work_dir: str, save_path = osp.join(work_dir, save_file) from_onnx( onnx_model, - osp.splitext(save_path)[0], + save_path, input_shapes=input_shapes, log_level=get_trt_log_level(), fp16_mode=final_params.get('fp16_mode', False), diff --git a/mmdeploy/backend/tensorrt/utils.py b/mmdeploy/backend/tensorrt/utils.py index de60aa1124..db283f8cf0 100644 --- a/mmdeploy/backend/tensorrt/utils.py +++ b/mmdeploy/backend/tensorrt/utils.py @@ -96,7 +96,7 @@ def cmd_result(txt: str): def from_onnx(onnx_model: Union[str, onnx.ModelProto], - output_file_prefix: str, + output_path: str, input_shapes: Dict[str, Sequence[int]], max_workspace_size: int = 0, fp16_mode: bool = False, @@ -109,7 +109,7 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], Args: onnx_model (str or onnx.ModelProto): Input onnx model to convert from. - output_file_prefix (str): The path to save the output ncnn file. + output_path (str): The path to save the output engine file. input_shapes (Dict[str, Sequence[int]]): The min/opt/max shape of each input. max_workspace_size (int): To set max workspace size of TensorRT engine. @@ -118,7 +118,7 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], Defaults to `False`. int8_mode (bool): Specifying whether to enable int8 mode. Defaults to `False`. - int8_param (dict): A dict of parameter int8 mode. Defaults to `None`. + int8_param (dict): A dict of int8 mode parameters. Defaults to `None`. device_id (int): Choice the device to create engine. Defaults to `0`. log_level (trt.Logger.Severity): The log level of TensorRT. Defaults to `trt.Logger.ERROR`. @@ -220,13 +220,11 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], if int8_mode: if not getattr(builder, 'platform_has_fast_int8', True): logger.warning('Platform does not has fast native int8.') - from .calib_utils import HDF5Calibrator + from .calib_utils import IteratorCalibrator config.set_flag(trt.BuilderFlag.INT8) assert int8_param is not None - config.int8_calibrator = HDF5Calibrator( + config.int8_calibrator = IteratorCalibrator( int8_param['calib_file'], - input_shapes, - model_type=int8_param['model_type'], device_id=device_id, algorithm=int8_param.get( 'algorithm', trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2)) @@ -242,7 +240,7 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], assert engine is not None, 'Failed to create TensorRT engine' - save(engine, output_file_prefix + '.engine') + save(engine, output_path) return engine diff --git a/mmdeploy/backend/tensorrt/wrapper.py b/mmdeploy/backend/tensorrt/wrapper.py index bd4034eacd..c0b3c74b13 100644 --- a/mmdeploy/backend/tensorrt/wrapper.py +++ b/mmdeploy/backend/tensorrt/wrapper.py @@ -67,7 +67,7 @@ class TRTWrapper(BaseWrapper): output_names should be the same as onnx model. Examples: - >>> from mmdeploy.backend.tensorrt import TRTWrapper + >>> from mmdeploy.backend.tensorrt.wrapper import TRTWrapper >>> engine_file = 'resnet.engine' >>> model = TRTWrapper(engine_file) >>> inputs = dict(input=torch.randn(1, 3, 224, 224)) diff --git a/mmdeploy/backend/torchscript/__init__.py b/mmdeploy/backend/torchscript/__init__.py index 7cb34c4a68..be25c5594a 100644 --- a/mmdeploy/backend/torchscript/__init__.py +++ b/mmdeploy/backend/torchscript/__init__.py @@ -6,10 +6,8 @@ _BackendManager = TorchScriptManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param __all__ = ['get_ops_path', 'ops_available', 'TorchScriptManager'] - -if is_available(): - from .wrapper import TorchscriptWrapper - - __all__ += ['TorchscriptWrapper'] diff --git a/mmdeploy/backend/torchscript/backend_manager.py b/mmdeploy/backend/torchscript/backend_manager.py index 91572f48d3..22e5715f09 100644 --- a/mmdeploy/backend/torchscript/backend_manager.py +++ b/mmdeploy/backend/torchscript/backend_manager.py @@ -1,38 +1,44 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging -from typing import Any, Callable, Optional, Sequence +import os.path as osp +import shutil +from dataclasses import dataclass +from typing import Any, Callable, List, Optional, Sequence -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.torchscript import TorchScriptParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor) -@BACKEND_MANAGERS.register('torchscript') -class TorchScriptManager(BaseBackendManager): +# We name the name `TorchJIT` to distinguish `Torchscript` as IR. +@dataclass +class TorchJITParam(BaseBackendParam): + """TorchJIT backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + input_names (List[str]): Names of the inputs. + output_names (List[str]): Names of the outputs. + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.pth') + + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str), ('Expect string work_dir, ' + f'got {self.work_dir}') + assert isinstance(self.file_name, str), ('Expect string file_name, ' + f'got {self.file_name}') + return osp.join(self.work_dir, self.file_name) + + +_BackendParam = TorchJITParam - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import TorchscriptWrapper - return TorchscriptWrapper( - model=backend_files[0], - input_names=input_names, - output_names=output_names) + +@BACKEND_MANAGERS.register( + 'torchscript', param=_BackendParam, ir_param=TorchScriptParam) +class TorchScriptManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -84,21 +90,87 @@ def check_env(cls, log_callback: Callable = lambda _: _) -> str: return info @classmethod - def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + def to_backend(cls, torhscript_path: str, save_path: str): """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be save in this directory. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. + torhscript_path (str): The intermediate representation files. + save_path (str): The save path of onnx path. Returns: Sequence[str]: Backend files. """ - return ir_files + if osp.abspath(save_path) != osp.abspath(torhscript_path): + shutil.copy(torhscript_path, save_path) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: BaseBackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + save_path = osp.join(param.work_dir, param.file_name) + cls.to_backend(ir_model, save_path) + + @classmethod + def build_wrapper(cls, + model_path: str, + input_names: Optional[Sequence[str]] = None, + output_names: Optional[Sequence[str]] = None): + """Build the wrapper for the backend model. + + Args: + model_path (str): torchscript model path. + input_names (Optional[Sequence[str]], optional): input names. + Defaults to None. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + """ + from .wrapper import TorchscriptWrapper + return TorchscriptWrapper( + model=model_path, + input_names=input_names, + output_names=output_names) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (_BackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + input_names = param.input_names + output_names = param.output_names + return cls.build_wrapper( + model_path, input_names=input_names, output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: List[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + from mmdeploy.utils import get_ir_config + ir_config = get_ir_config(config) + input_names = ir_config.get('input_names', []) + output_names = ir_config.get('output_names', []) + kwargs.update(dict(input_names=input_names, output_names=output_names)) + return _BackendParam( + work_dir=work_dir, file_name=backend_files[0], **kwargs) diff --git a/mmdeploy/backend/tvm/__init__.py b/mmdeploy/backend/tvm/__init__.py index 8389b527a4..ada9ea9199 100644 --- a/mmdeploy/backend/tvm/__init__.py +++ b/mmdeploy/backend/tvm/__init__.py @@ -1,38 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -import sys - -from .backend_manager import TVMManager +from .backend_manager import TVMManager, TVMParam, get_library_ext _BackendManager = TVMManager is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param - -def get_library_ext() -> str: - """Get the extension of the library. - - Returns: - str: The extension name - """ - platform = sys.platform.lower() - if platform == 'win32' or platform == 'cygwin': - return '.dll' - elif platform == 'linux' or platform == 'darwin' or platform == 'freebsd': - return '.so' - - -__all__ = ['TVMManager'] - -if is_available(): - from .onnx2tvm import from_onnx - from .quantize import HDF5Dataset - from .tuner import build_tvm_tuner - - __all__ += ['from_onnx', 'build_tvm_tuner', 'HDF5Dataset'] - - try: - # import wrapper if pytorch is available - from .wrapper import TVMWrapper - __all__ += ['TVMWrapper'] - except Exception: - pass +__all__ = ['TVMParam', 'TVMManager', 'get_library_ext'] diff --git a/mmdeploy/backend/tvm/backend_manager.py b/mmdeploy/backend/tvm/backend_manager.py index f9e99f18f8..a236bcce3e 100644 --- a/mmdeploy/backend/tvm/backend_manager.py +++ b/mmdeploy/backend/tvm/backend_manager.py @@ -1,41 +1,117 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib import os.path as osp -from typing import Any, Optional, Sequence +import re +import sys +from argparse import Action, ArgumentParser +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Optional, Sequence, Union -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + FileNameDescriptor, get_obj_by_qualname, + import_custom_modules) -@BACKEND_MANAGERS.register('tvm') -class TVMManager(BaseBackendManager): +def get_library_ext() -> str: + """Get the extension of the library. + + Returns: + str: The extension name + """ + platform = sys.platform.lower() + if platform == 'win32' or platform == 'cygwin': + return '.dll' + elif platform == 'linux' or platform == 'darwin' or platform == 'freebsd': + return '.so' + + +class DTypeAction(Action): + """dtype argparse action.""" + + def __call__(self, parser, namespace, values, option_string=None): + """call action.""" + args = values + if isinstance(args, str): + args = [args] + + pattern = r'^[^\S\n]*((?P\w+):)?(?P\w+)$' + ret = dict() + for arg in args: + arg = [i.strip() for i in arg.split(',')] + for single_arg in arg: + if len(single_arg) == 0: + continue + m = re.match(pattern, single_arg) + if m is None: + raise ValueError(f'Can not parse shape: {single_arg}') + input_name = m.group('input_name') + dtype = m.group('dtype') + if input_name in ret: + raise NameError( + f'shape of `{input_name}` has been assigned' + 'more than once.') + ret[input_name] = dtype + + setattr(namespace, self.dest, ret) + + +@dataclass +class TVMParam(BaseBackendParam): + """TVM backend parameters. + + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + vm_name (str): Serialized vm file. If not given, + vm_name would be he same as file_name with postfix `.vm` + use_vm (bool): Enable tvm virtual machine runtime. Defaults to False. + input_shapes (ShapeType): The Default shape of the inputs. + output_names (List[str]): Names of the outputs. + dtypes (Dict[str, str]): The input data types. + tuner (Optional[Union[TVMTunerBase, Dict]], optional): The tuner + config. Defaults to None. + qconfig (QConfig): `relay.quantize.QConfig` instance. + quanti_data (Any): Calibration dataset. Iterable object of + `Dict[str, ndarray]` + device (str): Device used to perform inference. + """ + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix=get_library_ext()) + vm_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.vm', base_name='file_name') + use_vm: bool = False + dtypes: Dict[str, str] = None + tuner: Any = None + qconfig: Any = None + device: str = 'llvm' + + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str) + assert isinstance(self.file_name, str) + param_file_path = osp.join(self.work_dir, self.file_name) + assert isinstance(self.vm_name, str) + algorithm_file_path = osp.join(self.work_dir, self.vm_name) + return param_file_path, algorithm_file_path @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + def add_argument(cls, parser: ArgumentParser, name: str, dtype: Any, + default: Any, desc: str): + arg_name = f'--{name.replace("_", "-")}' + if name == 'dtypes': + parser.add_argument( + arg_name, action=DTypeAction, nargs='+', help=desc) + else: + return super().add_argument(parser, name, dtype, default, desc) - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import TVMWrapper - bytecode = None if len(backend_files) <= 1 else backend_files[1] - return TVMWrapper( - backend_files[0], - bytecode=bytecode, - output_names=output_names, - device=device) + +_BackendParam = TVMParam + + +@BACKEND_MANAGERS.register('tvm', param=TVMParam, ir_param=ONNXParam) +class TVMManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -65,71 +141,243 @@ def get_version(cls) -> str: @classmethod def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + onnx_file: str, + output_file: str, + use_vm: bool = False, + vm_file: str = '', + input_shapes: Optional[Dict] = None, + dtypes: Union[str, Dict] = 'float32', + tuner: Optional[Union[Any, Dict, str]] = None, + qconfig: Optional[Union[Any, Dict, str]] = None, + dataset: Optional[Iterable] = None, + device: str = 'llvm') -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. + onnx_file (str): The intermediate representation files. + output_file (str): output library path + use_vm (bool, optional): Enable tvm virtual machine runtime. + Defaults to False. + vm_file (str, optional): output bytecode path for virtual + machine. Defaults to ''. + input_shapes (Optional[Dict], optional): The input shape + dictionary. Defaults to None. + dtypes (Union[str, Dict], optional): The input data type + dictionary. Defaults to 'float32'. + tuner (Optional[Union[TVMTunerBase, Dict]], optional): The tuner + config. Defaults to None. + qconfig (QConfig): `relay.quantize.QConfig` instance. + dataset (Any): Calibration dataset. Iterable object of + `Dict[str, ndarray]` + device (str): Device used to perform inference. Returns: Sequence[str]: Backend files. """ + from .onnx2tvm import from_onnx - import copy + # process dtypes + if isinstance(dtypes, Dict) and len(dtypes) == 1 and None in dtypes: + dtypes = dtypes[None] - from mmdeploy.apis.tvm import get_library_ext - from mmdeploy.utils import (get_calib_filename, get_model_inputs, - get_partition_config) - from .onnx2tvm import from_onnx - model_inputs = get_model_inputs(deploy_cfg) + # process tuner + if isinstance(tuner, str): + tuner = get_obj_by_qualname(tuner) + # process qconfig + if isinstance(qconfig, str): + qconfig = get_obj_by_qualname(qconfig) + + # process device if device.startswith('cuda'): - target = 'cuda' + device = 'cuda' else: - target = 'llvm' - - lib_ext = get_library_ext() - - tvm_files = [] - for model_id, onnx_path in enumerate(ir_files): - model_input = copy.deepcopy(model_inputs[model_id]) - use_vm = model_input.get('use_vm', False) - if 'target' not in model_input['tuner']: - model_input['tuner']['target'] = target - lib_path = osp.splitext(onnx_path)[0] + lib_ext - code_path = osp.splitext( - onnx_path)[0] + '.code' if use_vm else None - model_input['output_file'] = lib_path - model_input['onnx_model'] = onnx_path - model_input['bytecode_file'] = code_path - - # create calibration dataset - if 'qconfig' in model_input: - from .quantize import HDF5Dataset - calib_filename = get_calib_filename(deploy_cfg) - calib_path = osp.join(work_dir, calib_filename) - partition_cfgs = get_partition_config(deploy_cfg) - onnx_name = osp.splitext(osp.split(onnx_path)[1])[0] - partition_type = 'end2end' if partition_cfgs is None \ - else onnx_name - dataset = HDF5Dataset( - calib_path, - model_input['shape'], - model_type=partition_type, - device=target) - model_input['dataset'] = dataset() - - from_onnx(**model_input) - - tvm_files += [lib_path, code_path] - - return tvm_files + device = 'llvm' + + from_onnx( + onnx_file, + output_file=output_file, + use_vm=use_vm, + bytecode_file=vm_file, + shape=input_shapes, + dtype=dtypes, + tuner=tuner, + qconfig=qconfig, + dataset=dataset, + device=device) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + model_path = osp.join(param.work_dir, param.file_name) + assert isinstance(param.vm_name, str) + vm_path = osp.join(param.work_dir, param.vm_name) + + cls.to_backend( + ir_model, + model_path, + use_vm=param.use_vm, + vm_file=vm_path, + input_shapes=param.input_shapes, + dtypes=param.dtypes, + tuner=param.tuner, + qconfig=param.qconfig, + dataset=param.quanti_data, + device=param.device) + + @classmethod + def build_wrapper( + cls, + lib_file: str, + output_names: Optional[Sequence[str]], + vm_file: str = None, + device: str = 'cpu', + ): + """Build the wrapper for the backend model. + + Args: + lib_file (str): generated library path + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + vm_file (str, optional): output bytecode path for virtual + machine. Defaults to ''. + device (str, optional): The device info. Defaults to 'cpu'. + """ + from .wrapper import TVMWrapper + if vm_file is not None and not osp.exists(vm_file): + vm_file = None + return TVMWrapper( + lib_file, + output_names=output_names, + bytecode=vm_file, + device=device) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_path, vm_path = param.get_model_files() + output_names = param.output_names + device = param.device + return cls.build_wrapper( + model_path, + output_names=output_names, + vm_file=vm_path, + device=device) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + from mmdeploy.utils import (get_calib_config, get_ir_config, + get_model_inputs) + + # get output names + ir_config = get_ir_config(config) + output_names = ir_config['output_names'] + kwargs.setdefault('output_names', output_names) + + # get tvm param + model_inputs = get_model_inputs(config) + shape = model_inputs[0]['shape'] + dtype = model_inputs[0].get('dtype', 'float32') + tuner = model_inputs[0].get('tuner', dict(type='DefaultTuner')) + + kwargs.setdefault('work_dir', work_dir) + kwargs.setdefault('input_shapes', shape) + kwargs.setdefault('dtypes', dtype) + kwargs.setdefault('tuner', tuner) + + qconfig = model_inputs[0].get('qconfig', None) + if qconfig is not None: + from ..base import create_h5pydata_generator + kwargs.setdefault('qconfig', qconfig) + calib_config = get_calib_config(config) + assert 'calib_file' in calib_config + calib_path = osp.join(work_dir, calib_config['calib_file']) + kwargs.setdefault('quanti_data', + create_h5pydata_generator(calib_path, shape)) + + backend_files = [] if backend_files is None else backend_files + if len(backend_files) > 0: + kwargs['file_name'] = backend_files[0] + if len(backend_files) > 1: + kwargs['vm_name'] = backend_files[1] + return _BackendParam(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Custom module path.') + + parsed_args = parser.parse_args(args) + yield parsed_args + import_custom_modules(parsed_args.custom_modules) + + # perform command + command = parsed_args._command + + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + vm_name=parsed_args.vm_name, + use_vm=parsed_args.use_vm, + input_shapes=parsed_args.input_shapes, + output_names=parsed_args.output_names, + dtypes=parsed_args.dtypes, + tuner=parsed_args.tuner, + qconfig=parsed_args.qconfig, + quanti_data=parsed_args.quanti_data, + device=parsed_args.device) + + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/tvm/onnx2tvm.py b/mmdeploy/backend/tvm/onnx2tvm.py index 54ca33e2c8..c248ab9831 100644 --- a/mmdeploy/backend/tvm/onnx2tvm.py +++ b/mmdeploy/backend/tvm/onnx2tvm.py @@ -1,5 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -from typing import Callable, Dict, Optional, Union +from typing import Dict, Iterable, Optional, Union import onnx from tvm.relay.frontend import from_onnx as relay_from_onnx @@ -20,7 +20,8 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], dtype: Union[str, Dict] = 'float32', tuner: Optional[Union[TVMTunerBase, Dict]] = None, qconfig: Optional[Union[QConfig, Dict]] = None, - dataset: Optional[Callable] = None): + dataset: Optional[Iterable] = None, + device: str = 'llvm'): """Convert ONNX model to tvm lib. Args: @@ -30,12 +31,16 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], Defaults to False. bytecode_file (str, optional): output bytecode path for virtual machine. Defaults to ''. - shape (Optional[Dict], optional): The input shape directory. Defaults + shape (Optional[Dict], optional): The input shape dictionary. Defaults to None. dtype (Union[str, Dict], optional): The input data type dictionary. Defaults to 'float32'. tuner (Optional[Union[TVMTunerBase, Dict]], optional): The tuner config. Defaults to None. + qconfig (QConfig): `relay.quantize.QConfig` instance. + dataset (Any): Calibration dataset. Iterable object of + `Dict[str, ndarray]` + device (str): Device use to create quantization dataset. Return: lib: The converted tvm lib @@ -43,7 +48,7 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], None if use_vm==False. Examples: - >>> from mmdeploy.backend.tvm import from_onnx + >>> from mmdeploy.backend.tvm.onnx2tvm import from_onnx >>> onnx_path = 'model.onnx' >>> output_file = 'model.so' >>> shape = {'input':[1,3,224,224]} @@ -73,16 +78,21 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto], qconfig = create_qconfig(**qconfig) with qconfig: - mod = quantize(mod, params, dataset) + from .quantize import IteratorDataset + iter_dataset = IteratorDataset(dataset, device) + mod = quantize(mod, params, iter_dataset()) if tuner is None: # use default tuner - tuner = dict(type='DefaultTuner', target=Target('llvm')) + tuner = dict(type='DefaultTuner', target=Target(device)) if not issubclass(type(tuner), TVMTunerBase): tuner['use_vm'] = use_vm + tuner['target'] = Target(device) tuner = build_tvm_tuner(tuner) + tuner._target = Target(device) + logger.info(f'Tuning with {type(tuner).__name__} .') tuner.tune(mod, params) lib = tuner.build(mod, params) diff --git a/mmdeploy/backend/tvm/quantize.py b/mmdeploy/backend/tvm/quantize.py index ba69e918bf..3dffeb9a66 100644 --- a/mmdeploy/backend/tvm/quantize.py +++ b/mmdeploy/backend/tvm/quantize.py @@ -1,45 +1,26 @@ # Copyright (c) OpenMMLab. All rights reserved. -from typing import Any, Dict, Sequence, Union +from typing import Iterable -import numpy as np import tvm from tvm.runtime.ndarray import array -class HDF5Dataset: +class IteratorDataset: """HDF5 dataset. Args: - calib_file (str | h5py.File): Input calibration file. - input_shapes (Dict[str, Sequence[int]]): The shape of - each input. - model_type (str): Input model type, defaults to 'end2end'. + dataset (Iterable): Iterable dataset object. device (str): Device type, default to llvm. """ def __init__( self, - calib_file: Union[str, Any], - input_shapes: Dict[str, Sequence[int]], - model_type: str = 'end2end', + dataset: Iterable, device: str = 'llvm', ) -> None: - import h5py - if isinstance(calib_file, str): - calib_file = h5py.File(calib_file, mode='r') - assert 'calib_data' in calib_file - calib_data = calib_file['calib_data'] - assert model_type in calib_data - calib_data = calib_data[model_type] - - self.calib_file = calib_file - self.calib_data = calib_data - self.device = device - self.input_shapes = input_shapes - - first_input_group = calib_data[list(calib_data.keys())[0]] - self.dataset_length = len(first_input_group) + self._dataset = dataset + self._device = device def __call__(self): """Create dataset generator. @@ -47,27 +28,18 @@ def __call__(self): Yields: Iterator[Any]: data in the dataset """ - for idx in range(self.dataset_length): - + for data in self._dataset: ret = dict() - for name, opt_shape in self.input_shapes.items(): - input_group = self.calib_data[name] - data_np = input_group[str(idx)][...].astype(np.float32) - - data_shape = data_np.shape - - # tile the input data - reps = [ - int(np.ceil(opt_s / data_s)) - for opt_s, data_s in zip(opt_shape, data_shape) - ] - - data_np = np.tile(data_np, reps) - - slice_list = tuple(slice(0, end) for end in opt_shape) - data_np = data_np[slice_list] - - data_nd = array(data_np, tvm.device(self.device)) + for name, data_np in data.items(): + # cast back to numpy + try: + import torch + if isinstance(data_np, torch.Tensor): + data_np = data_np.detach().cpu().numpy() + except Exception: + pass + + # tvm array + ret[name] = array(data_np, tvm.device(self._device)) - ret[name] = data_nd yield ret diff --git a/mmdeploy/backend/tvm/wrapper.py b/mmdeploy/backend/tvm/wrapper.py index e77d02c9ef..3f076f957a 100644 --- a/mmdeploy/backend/tvm/wrapper.py +++ b/mmdeploy/backend/tvm/wrapper.py @@ -24,7 +24,7 @@ class TVMWrapper(BaseWrapper): Examples: - >>> from mmdeploy.backend.tvm import TVMWrapper + >>> from mmdeploy.backend.tvm.wrapper import TVMWrapper >>> lib_file = 'resnet.so' >>> model = TVMWrapper(lib_file, ['output']) >>> inputs = dict(input=torch.randn(1, 3, 224, 224)) diff --git a/mmdeploy/backend/vacc/__init__.py b/mmdeploy/backend/vacc/__init__.py index ccbf8827ae..ff11e422e7 100644 --- a/mmdeploy/backend/vacc/__init__.py +++ b/mmdeploy/backend/vacc/__init__.py @@ -1,18 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .backend_manager import VACCManager -from .onnx2vacc import from_onnx +from .backend_manager import VACCManager, VACCParam _BackendManager = VACCManager - is_available = _BackendManager.is_available build_wrapper = _BackendManager.build_wrapper +build_wrapper_from_param = _BackendManager.build_wrapper_from_param +to_backend = _BackendManager.to_backend +to_backend_from_param = _BackendManager.to_backend_from_param -__all__ = ['VACCManager', 'from_onnx'] - -if is_available(): - try: - from .wrapper import VACCWrapper - - __all__ += ['VACCWrapper'] - except Exception: - pass +__all__ = ['VACCParam', 'VACCManager'] diff --git a/mmdeploy/backend/vacc/backend_manager.py b/mmdeploy/backend/vacc/backend_manager.py index 813162074b..b3eebef82f 100644 --- a/mmdeploy/backend/vacc/backend_manager.py +++ b/mmdeploy/backend/vacc/backend_manager.py @@ -1,53 +1,63 @@ # Copyright (c) OpenMMLab. All rights reserved. -import logging +import contextlib +import os.path as osp import sys -from typing import Any, Callable, Optional, Sequence +from argparse import ArgumentParser +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Sequence, Union -from mmdeploy.utils import get_common_config, get_model_inputs, get_root_logger -from ..base import BACKEND_MANAGERS, BaseBackendManager +from mmdeploy.ir.onnx import ONNXParam +from mmdeploy.utils import get_root_logger +from ..base import (BACKEND_MANAGERS, BaseBackendManager, BaseBackendParam, + get_obj_by_qualname, import_custom_modules) -@BACKEND_MANAGERS.register('vacc') -class VACCManager(BaseBackendManager): +@dataclass +class VACCParam(BaseBackendParam): + """VACC backend parameters. - @classmethod - def build_wrapper(cls, - backend_files: Sequence[str], - device: str = 'cpu', - input_names: Optional[Sequence[str]] = None, - output_names: Optional[Sequence[str]] = None, - deploy_cfg: Optional[Any] = None, - **kwargs): - """Build the wrapper for the backend model. + Args: + work_dir (str): The working directory. + file_name (str): File name of the serialized model. Postfix will be + added automatically. + quant_mode (str): quantization mode, choice between ['fp16', 'int8'] + input_shapes (ShapeType): The Default shape of the inputs. + output_names (List[str]): Names of the outputs. + calib_num (int): Max numbers of calibration data. + qconfig (Dict): Dictionary arguments feed to vacc.qconfig. + Or qualname to the dict. + quanti_data (Any): Calibration dataset. Iterable object of + `Dict[str, ndarray]` + data_transmode (int): `tvm.build_config` arguments. + cluster_mode (int): `tvm.build_config` arguments. + vdsp_params_info (str|Dict): vdsp parameters file or qualname of the + parameters dictionary. + """ + quant_mode: str = 'fp16' + calib_num: int = 1000 + qconfig: Union[str, Dict] = field(default_factory=dict) + data_transmode: int = 1 + cluster_mode: int = 0 + vdsp_params_info: Union[str, Dict] = None - Args: - backend_files (Sequence[str]): Backend files. - device (str, optional): The device info. Defaults to 'cpu'. - input_names (Optional[Sequence[str]], optional): input names. - Defaults to None. - output_names (Optional[Sequence[str]], optional): output names. - Defaults to None. - deploy_cfg (Optional[Any], optional): The deploy config. Defaults - to None. - """ - from .wrapper import VACCWrapper + def get_model_files(self) -> str: + """get the model files.""" + assert isinstance(self.work_dir, str) + assert isinstance(self.file_name, str) + save_dir = '-'.join([self.file_name, self.quant_mode]) + name = osp.split(self.file_name)[1] + model_prefix = osp.join(self.work_dir, save_dir, name) + return [ + model_prefix + '.so', model_prefix + '.json', + model_prefix + '.params' + ] - # For unittest deploy_config will not pass into _build_wrapper - # function. - try: - common_cfg = get_common_config(deploy_cfg) - vdsp_params_info = common_cfg['vdsp_params_info'] +_BackendParam = VACCParam - return VACCWrapper( - lib_file=backend_files[0], - graph_file=backend_files[1], - param_file=backend_files[2], - vdsp_params_info=vdsp_params_info, - output_names=output_names) - except Exception: - print('Build model process success, wrapper process stopped') - exit(1) + +@BACKEND_MANAGERS.register('vacc', param=VACCParam, ir_param=ONNXParam) +class VACCManager(BaseBackendManager): @classmethod def is_available(cls, with_custom_ops: bool = False) -> bool: @@ -98,31 +108,30 @@ def check_env(cls, log_callback: Callable = lambda _: _) -> str: return info @classmethod - def to_backend(cls, - ir_files: Sequence[str], - work_dir: str, - deploy_cfg: Any, - log_level: int = logging.INFO, - device: str = 'cpu', - **kwargs) -> Sequence[str]: + def to_backend( + cls, + onnx_model: str, + output_path: str, + model_name: str, + input_shapes: Dict[str, Sequence], + quant_mode: str = 'fp16', + calib_num: int = 1000, + qconfig: Optional[Dict] = None, + data_transmode: int = 1, + cluster_mode: int = 0, + ) -> Sequence[str]: """Convert intermediate representation to given backend. Args: - ir_files (Sequence[str]): The intermediate representation files. - work_dir (str): The work directory, backend files and logs should - be saved in this directory. - deploy_cfg (Any): The deploy config. - log_level (int, optional): The log level. Defaults to logging.INFO. - device (str, optional): The device type. Defaults to 'cpu'. - Returns: - Sequence[str]: Backend files. + onnx_model (str): Input onnx model. + output_path (str): File path to save VACC model. + model_name (str): model name. + input_shapes (ShapeType): The Default shape of the inputs. + qconfig (Dict): Dictionary arguments feed to vacc.qconfig. """ logger = get_root_logger() - import copy - - from . import is_available - if not is_available(): + if not cls.is_available(): logger.error( 'vacc and tvm support is not available, please make sure:\n' '1) `vacc/python` and `tvm/python` existed in `PYTHONPATH`\n' @@ -130,15 +139,203 @@ def to_backend(cls, sys.exit(1) from .onnx2vacc import from_onnx - model_inputs = get_model_inputs(deploy_cfg) + + if isinstance(qconfig, str): + qconfig = get_obj_by_qualname(qconfig) + assert isinstance(qconfig, Dict) + + from_onnx( + onnx_model, + output_path=output_path, + model_name=model_name, + input_shapes=input_shapes, + quant_mode=quant_mode, + calib_num=calib_num, + qconfig=qconfig, + data_transmode=data_transmode, + cluster_mode=cluster_mode) + + @classmethod + def to_backend_from_param(cls, ir_model: str, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + ir_model (str): The ir model path to perform the export. + param (BaseBackendParam): Packed backend parameter. + """ + assert isinstance(param, _BackendParam) + assert isinstance(param.work_dir, str) + assert isinstance(param.file_name, str) + + cls.to_backend( + ir_model, + output_path=param.work_dir, + model_name=param.file_name, + input_shapes=param.input_shapes, + quant_mode=param.quant_mode, + calib_num=param.calib_num, + qconfig=param.qconfig, + data_transmode=param.data_transmode, + cluster_mode=param.cluster_mode) + + @classmethod + def build_wrapper(cls, + lib_file: str, + graph_file: str, + param_file: str, + vdsp_params_info: dict, + output_names: Optional[Sequence[str]] = None): + """Build the wrapper for the backend model. + + Args: + lib_file (str): Path of a model lib file. + graph_file (str): Path of a model graph file. + param_file (str): Path of a model param file. + vdsp_params_info_json (str): Path of a vdsp params info json file. + output_names (Optional[Sequence[str]], optional): output names. + Defaults to None. + """ + from .wrapper import VACCWrapper + + # For unittest deploy_config will not pass into _build_wrapper + # function. + + if isinstance(vdsp_params_info, + str) and not osp.exists(vdsp_params_info): + vdsp_params_info = get_obj_by_qualname(vdsp_params_info) + + try: + return VACCWrapper( + lib_file=lib_file, + graph_file=graph_file, + param_file=param_file, + vdsp_params_info=vdsp_params_info, + output_names=output_names) + except Exception as e: + print(f'failed with error: {e}') + print('Build model process success, wrapper process stopped') + exit(1) + + @classmethod + def build_wrapper_from_param(cls, param: _BackendParam): + """Export to backend with packed backend parameter. + + Args: + param (BaseBackendParam): Packed backend parameter. + """ + model_paths = param.get_model_files() + output_names = param.output_names + vdsp_params_info = param.vdsp_params_info + return cls.build_wrapper( + *model_paths, + vdsp_params_info=vdsp_params_info, + output_names=output_names) + + @classmethod + def build_param_from_config(cls, + config: Any, + work_dir: str, + backend_files: Sequence[str] = None, + **kwargs) -> _BackendParam: + """Build param from deploy config. + + Args: + config (Any): The deploy config. + work_dir (str): work directory of the parameters. + backend_files (List[str]): The backend files of the model. + + Returns: + BaseBackendParam: The packed backend parameter. + """ + + from mmdeploy.utils import (get_common_config, get_ir_config, + get_model_inputs) + + deploy_cfg = config + ir_cfg = get_ir_config(deploy_cfg) + model_inputs = get_model_inputs(deploy_cfg)[0] common_params = get_common_config(deploy_cfg) model_name = common_params['name'] + output_names = ir_cfg.get('output_names', None) + vdsp_params_info = common_params['vdsp_params_info'] + input_shapes = model_inputs.get('shape', None) + qconfig = model_inputs.get('qconfig', {}) + quant_mode = qconfig.pop('dtype', 'fp16') + calib_num = qconfig.pop('calib_num', 1000) + data_transmode = qconfig.pop('data_transmode', 1) + cluster_mode = qconfig.pop('cluster_mode', 1) + + kwargs.setdefault('output_names', output_names) + kwargs.setdefault('vdsp_params_info', vdsp_params_info) + kwargs.setdefault('input_shapes', input_shapes) + kwargs.setdefault('qconfig', qconfig) + kwargs.setdefault('quant_mode', quant_mode) + kwargs.setdefault('calib_num', calib_num) + kwargs.setdefault('data_transmode', data_transmode) + kwargs.setdefault('cluster_mode', cluster_mode) + + kwargs.setdefault('work_dir', work_dir) + if len(backend_files) == 1: + file_name = model_name + else: + lib_path = osp.join(work_dir, backend_files[0]) + lib_dir = osp.split(lib_path)[0] + file_name = lib_dir[:-len(quant_mode) - 1] + + kwargs.setdefault('file_name', file_name) + + return cls.build_param(**kwargs) + + @classmethod + @contextlib.contextmanager + def parse_args(cls, + parser: ArgumentParser, + args: Optional[List[str]] = None): + """Parse console arguments. + + Args: + parser (ArgumentParser): The parser used to parse arguments. + args (Optional[List[str]], optional): Arguments to be parsed. If + not given, arguments from console will be parsed. + """ + + # parse args + sub_parsers = parser.add_subparsers( + title='command', + description='Please select the command you want to perform.', + dest='_command') + + # export model + export_parser = sub_parsers.add_parser( + name='convert', help='convert model from ONNX model.') + export_parser.add_argument( + '--onnx-path', required=True, help='ONNX model path.') + _BackendParam.add_arguments(export_parser) + export_parser.add_argument( + '--custom-modules', + type=str, + nargs='*', + help='Custom module path.') + + parsed_args = parser.parse_args(args) + yield parsed_args + import_custom_modules(parsed_args.custom_modules) + + # perform command + command = parsed_args._command - backend_files = [] - for model_id, onnx_path in zip(range(len(ir_files)), ir_files): - model_input = copy.deepcopy(model_inputs[model_id]) - model_file = from_onnx(onnx_path, work_dir, model_input, - model_name) - backend_files += model_file + if command == 'convert': + # convert model + param = _BackendParam( + work_dir=parsed_args.work_dir, + file_name=parsed_args.file_name, + input_shapes=parsed_args.input_shapes, + output_names=parsed_args.output_names, + quant_mode=parsed_args.quant_mode, + calib_num=parsed_args.calib_num, + qconfig=parsed_args.qconfig, + data_transmode=parsed_args.data_transmode, + vdsp_params_info=parsed_args.vdsp_params_info, + cluster_mode=parsed_args.cluster_mode) - return backend_files + cls.to_backend_from_param(parsed_args.onnx_path, param) diff --git a/mmdeploy/backend/vacc/onnx2vacc.py b/mmdeploy/backend/vacc/onnx2vacc.py index d34b5937a6..950896bbb0 100644 --- a/mmdeploy/backend/vacc/onnx2vacc.py +++ b/mmdeploy/backend/vacc/onnx2vacc.py @@ -1,30 +1,43 @@ # Copyright (c) OpenMMLab. All rights reserved. import os import os.path as osp - -import onnx -import tvm -import tvm.relay as relay -from vacc import quantize - - -def from_onnx(onnx_model: str, output_path: str, model_input: dict, - model_name: str, **kwargs): +from typing import Dict, Optional, Sequence + + +def from_onnx( + onnx_model: str, + output_path: str, + model_name: str, + input_shapes: Dict[str, Sequence], + quant_mode: str = 'fp16', + calib_num: int = 1000, + qconfig: Optional[Dict] = None, + data_transmode: int = 1, + cluster_mode: int = 0, +): """Convert ONNX to VACC. Args: onnx_model (str): Input onnx model. output_path (str): File path to save VACC model. - model_input (dict): model input config. model_name (str): model name. + input_shapes (ShapeType): The Default shape of the inputs. + quant_mode (str): quantization mode, choice between ['fp16', 'int8'] + calib_num (int): Max numbers of calibration data. + qconfig (Dict): Dictionary arguments feed to vacc.qconfig. + data_transmode (int): `tvm.build_config` arguments. + cluster_mode (int): `tvm.build_config` arguments. """ + import onnx + import tvm + import tvm.relay as relay + from vacc import quantize target = tvm.target.vacc() - quant_mode = model_input.get('qconfig', {}).get('dtype', 'fp16') assert quant_mode in ['int8', 'fp16'], quant_mode + ' not support now' - shape_dict = model_input['shape'] + shape_dict = input_shapes mod, params = relay.frontend.from_onnx(onnx.load(onnx_model), shape_dict) func = mod['main'] @@ -40,23 +53,19 @@ def from_onnx(onnx_model: str, output_path: str, model_input: dict, index = list(range(len(data))) random.shuffle(index) - calib_num = model_input.get('qconfig', {}).get('calib_num', 1000) for i in index[:calib_num]: calib_data.append({ list(shape_dict.keys())[0]: tvm.nd.array(data[str(i)][:].astype('float32')) }) + if qconfig is None: + qconfig = dict() with quantize.qconfig( - calibrate_mode=model_input.get('qconfig', - {}).get('calibrate_mode', - 'percentile'), - skip_conv_layers=model_input.get('qconfig', {}).get( - 'skip_conv_layers', []), - weight_scale=model_input.get('qconfig', - {}).get('weight_scale', 'max'), - quantize_per_channel=model_input.get('qconfig', {}).get( - 'per_channel', False)): + calibrate_mode=qconfig.get('calibrate_mode', 'percentile'), + skip_conv_layers=qconfig.get('skip_conv_layers', []), + weight_scale=qconfig.get('weight_scale', 'max'), + quantize_per_channel=qconfig.get('per_channel', False)): qmod = quantize.quantize(mod, params, calib_data) @@ -70,11 +79,9 @@ def from_onnx(onnx_model: str, output_path: str, model_input: dict, with tvm.build_config( data_type=data_type, - data_transport_mode=model_input.get('qconfig', - {}).get('data_transmode', 1), + data_transport_mode=data_transmode, mem_inplace=True, - cluster_mode=model_input.get('qconfig', {}).get('cluster_mode', - 0)): + cluster_mode=cluster_mode): with relay.build_config( opt_level=2, stream_mode=True, enable_float_to_half=True): graph, lib, params = relay.build( @@ -96,10 +103,5 @@ def from_onnx(onnx_model: str, output_path: str, model_input: dict, with open(param_path, 'wb') as f: f.write(relay.save_param_dict(params)) - assert osp.exists(os.path.join(output_root, - model_name + '.params')), 'onnx2vacc failed' - return [ - os.path.join(output_root, model_name + '.so'), - os.path.join(output_root, model_name + '.json'), - os.path.join(output_root, model_name + '.params') - ] + assert osp.exists(param_path), 'onnx2vacc failed' + return [libpath, graph_json_path, param_path] diff --git a/mmdeploy/backend/vacc/wrapper.py b/mmdeploy/backend/vacc/wrapper.py index a5c5314d91..79aeac6067 100644 --- a/mmdeploy/backend/vacc/wrapper.py +++ b/mmdeploy/backend/vacc/wrapper.py @@ -161,18 +161,24 @@ def __init__(self, } model_info_json = json.dumps(model_info) - with open(os.path.join(parent_path, 'model_info.json'), - 'w') as json_file: + model_info_path = os.path.join(parent_path, 'model_info.json') + with open(model_info_path, 'w') as json_file: json_file.write(model_info_json) - vdsp_params_info_json = json.dumps(vdsp_params_info) - with open(os.path.join(parent_path, 'vdsp_param_info.json'), - 'w') as json_file: - json_file.write(vdsp_params_info_json) - - self.model = VACCForward( - os.path.join(parent_path, 'model_info.json'), - os.path.join(parent_path, 'vdsp_param_info.json')) + if isinstance(vdsp_params_info, Dict): + vdsp_params_info_path = os.path.join(parent_path, + 'vdsp_param_info.json') + vdsp_params_info_json = json.dumps(vdsp_params_info) + with open(vdsp_params_info_path, 'w') as json_file: + json_file.write(vdsp_params_info_json) + elif isinstance(vdsp_params_info, str): + vdsp_params_info_path = vdsp_params_info + else: + raise TypeError( + 'Expect vdsp_param_info be .json file or parameter dictory' + f' but get {type(vdsp_params_info)}') + + self.model = VACCForward(model_info_path, vdsp_params_info_path) super().__init__(output_names) diff --git a/mmdeploy/codebase/base/backend_model.py b/mmdeploy/codebase/base/backend_model.py index 327621eff1..5280c73585 100644 --- a/mmdeploy/codebase/base/backend_model.py +++ b/mmdeploy/codebase/base/backend_model.py @@ -62,8 +62,15 @@ def _build_wrapper(backend: Backend, if backend_mgr is None: raise NotImplementedError( f'Unsupported backend type: {backend.value}') - return backend_mgr.build_wrapper(backend_files, device, input_names, - output_names, deploy_cfg, **kwargs) + + param = backend_mgr.build_param_from_config( + deploy_cfg, + work_dir='', + backend_files=backend_files, + input_names=input_names, + output_names=output_names, + device=device) + return backend_mgr.build_wrapper_from_param(param) def destroy(self): if hasattr(self, 'wrapper') and hasattr(self.wrapper, 'destroy'): diff --git a/mmdeploy/ir/base/__init__.py b/mmdeploy/ir/base/__init__.py new file mode 100644 index 0000000000..500be2a51f --- /dev/null +++ b/mmdeploy/ir/base/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .ir_manager import (IR_MANAGERS, BaseIRManager, BaseIRParam, + FileNameDescriptor) + +__all__ = ['IR_MANAGERS', 'BaseIRManager', 'BaseIRParam', 'FileNameDescriptor'] diff --git a/mmdeploy/ir/base/ir_manager.py b/mmdeploy/ir/base/ir_manager.py new file mode 100644 index 0000000000..68c8d6cbb9 --- /dev/null +++ b/mmdeploy/ir/base/ir_manager.py @@ -0,0 +1,188 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import importlib +import os.path as osp +from abc import ABCMeta +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + + +class FileNameDescriptor: + """File name descriptor.""" + + def __init__(self, *, default, postfix: str = '', base_name=None): + self._default = default + self._postfix = postfix + self._base_name = base_name + + def __set_name__(self, owner, name): + """set obj name.""" + self._name = '_' + name + + def __get__(self, obj, type): + """file name getter.""" + if obj is None: + return self._default + + # get . + ret = getattr(obj, self._name, self._default) + + # if . is None, try get name from base name + if ret is None and self._base_name is not None: + base_val = getattr(obj, self._base_name, None) + if base_val is not None: + name = osp.splitext(base_val)[0] + ret = name + self._postfix + return ret + + def __set__(self, obj, val): + """file name setter.""" + if val is not None and osp.splitext(val)[1] == '': + val = val + self._postfix + setattr(obj, self._name, val) + + +@dataclass +class BaseIRParam: + """Base ir param. + + Args: + args (Any): The arguments of the model. + work_dir (str): The working directory to save the output. + file_name (str): The file name of the output. postfix can be omitted. + input_names (List[str]): The names to assign to the input of the ir. + output_names (List[str]): The names to assign to the output of the ir. + dynamic_axes (Dict): Determine the dynamic axes of the inputs. It not + given, all axes will be static. + backend (str): The expected backend of the ir. + rewrite_context (Dict): Provide information to the rewriter. + """ + # latent fields + _manager = None + + # class fields + args: Any = None + work_dir: str = None + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='') + input_names: List[str] = None + output_names: List[str] = None + dynamic_axes: Dict = None + backend: str = 'default' + rewrite_context: Dict = field(default_factory=dict) + + @classmethod + def get_manager(cls): + """manager of the ir.""" + return cls._manager + + def check(self): + """check if the param is valid.""" + + +class BaseIRManager(metaclass=ABCMeta): + """Abstract interface of ir manager.""" + + build_param = BaseIRParam + + @classmethod + def export(cls, model: Any, *args, **kwargs): + """export model to ir.""" + raise NotImplementedError( + 'class method: `export` of ' + f'{cls.__qualname__} has not been implemented.') + + @classmethod + def export_from_param(cls, model, param: BaseIRParam): + """export model to ir by param.""" + raise NotImplementedError( + 'class method: `export_from_param` of ' + f'{cls.__qualname__} has not been implemented.') + + @classmethod + def is_available(cls) -> bool: + """check if the export tools is available.""" + raise NotImplementedError( + 'class method: `is_available` of ' + f'{cls.__qualname__} has not been implemented.') + + +class IRManagerRegistry: + """ir manager registry.""" + + def __init__(self): + self._module_dict = {} + + def register(self, + name: str, + enum_name: Optional[str] = None, + param: Any = None): + """register ir manager. + + Args: + name (str): name of the ir + enum_name (Optional[str], optional): enum name of the ir. + if not given, the upper case of name would be used. + """ + from mmdeploy.utils import get_root_logger + logger = get_root_logger() + + if enum_name is None: + enum_name = name.upper() + + def wrap_manager(cls): + + from mmdeploy.utils import IR + + if not hasattr(IR, enum_name): + from aenum import extend_enum + extend_enum(IR, enum_name, name) + logger.info(f'Registry new ir: {enum_name} = {name}.') + + if name in self._module_dict: + logger.info( + f'IR manager of `{name}` has already been registered.') + + self._module_dict[name] = cls + + cls.ir_name = name + cls.build_param = param + if param is not None: + param._manager = cls + + return cls + + return wrap_manager + + def find(self, name: str) -> BaseIRManager: + """Find the ir manager with name. + + Args: + name (str): ir name. + Returns: + BaseIRManager: ir manager of the given ir. + """ + # try import name if it exists in `mmdeploy.ir` + try: + importlib.import_module(f'mmdeploy.ir.{name}') + except Exception: + from mmdeploy.utils import get_root_logger + logger = get_root_logger() + logger.debug(f'can not find IR: {name} in `mmdeploy.ir`') + return self._module_dict.get(name, None) + + +IR_MANAGERS = IRManagerRegistry() + + +def get_ir_manager(name: str) -> BaseIRManager: + """Get ir manager. + + Args: + name (str): name of the ir. + Returns: + BaseIRManager: The ir manager of given name + """ + from enum import Enum + if isinstance(name, Enum): + name = name.value + return IR_MANAGERS.find(name) diff --git a/mmdeploy/ir/onnx/__init__.py b/mmdeploy/ir/onnx/__init__.py new file mode 100644 index 0000000000..ea413a87be --- /dev/null +++ b/mmdeploy/ir/onnx/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .ir_manager import ONNXManager, ONNXParam + +export = ONNXManager.export +export_from_param = ONNXManager.export_from_param +is_available = ONNXManager.is_available + +__all__ = ['ONNXManager', 'ONNXParam'] diff --git a/mmdeploy/ir/onnx/ir_manager.py b/mmdeploy/ir/onnx/ir_manager.py new file mode 100644 index 0000000000..73c8f4e84f --- /dev/null +++ b/mmdeploy/ir/onnx/ir_manager.py @@ -0,0 +1,171 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os.path as osp +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +from ..base import IR_MANAGERS, BaseIRManager, BaseIRParam, FileNameDescriptor + + +@dataclass +class ONNXParam(BaseIRParam): + """ONNX IR param. + + Args: + args (Any): The arguments of the model. + work_dir (str): The working directory to save the output. + file_name (str): The file name of the output. postfix can be omitted. + input_names (List[str]): The names to assign to the input of the ir. + output_names (List[str]): The names to assign to the output of the ir. + dynamic_axes (Dict): Determine the dynamic axes of the inputs. It not + given, all axes will be static. + do_constant_folding (bool): Perform constant folding to the exported + model. Default to True. + opset_version (int): The version of the ONNX opset. Default to 11. + backend (str): The expected backend of the ir. + rewrite_context (Dict): Provide information to the rewriter. + verbose (bool): Show detail log of ONNX export. + const_args (Any): The constant args of the model. + optimize (bool): Perform optimization. + """ + # class fields + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.onnx') + do_constant_folding: bool = True + opset_version: int = 11 + verbose: bool = False + const_args: Any = None + optimize: bool = True + + def check(self): + super().check() + import torch + assert isinstance( + self.args, + (torch.Tensor, Tuple, + Dict)), ('Expect args type: (torch.Tensor, Sequence, Dict),', + f' get type: {type(self.args)}.') + assert self.opset_version >= 7, 'opset version < 7 is not supported.' + + +@IR_MANAGERS.register('onnx', param=ONNXParam) +class ONNXManager(BaseIRManager): + """ONNX IR Manager.""" + + @classmethod + def export(cls, + model: Any, + args: Any, + output_path: str, + input_names: Optional[List[str]] = None, + output_names: Optional[List[str]] = None, + opset_version: int = 11, + dynamic_axes: Optional[Dict] = None, + backend: str = 'default', + rewrite_context: Optional[Dict] = None, + verbose: bool = False, + const_args: Optional[Dict] = None, + optimize: bool = True): + """export model to ONNX. + + Examples: + >>> from mmdeploy.ir.onnx import export + >>> + >>> model = create_model() + >>> args = get_input_tensor() + >>> + >>> export( + >>> model, + >>> args, + >>> 'place/to/save/model.onnx', + >>> backend='tensorrt', + >>> input_names=['input'], + >>> output_names=['output'], + >>> dynamic_axes={'input': { + >>> 0: 'batch', + >>> 2: 'height', + >>> 3: 'width' + >>> }}) + + Args: + model (Any): Exportable PyTorch Model + args (Any): Arguments are used to trace the graph. + output_path (str): The output path. + input_names (List[str], optional): The name of the input in + the graph. Defaults to None. + output_names (List[str], optional): The name of the output + in the graph. Defaults to None. + opset_version (int): The ONNX opset version. Defaults to 11. + dynamic_axes (Dict, optional): Dynamic axes of each inputs. + If not given, all inputs share the fixed shapes of the args. + verbose (bool): Show detail export logs. Defaults to False. + const_args (Dict, optional): The non-exported inputs of the model. + rewrite_context (Dict, optional): The information used by + the rewriter. + optimize (bool): Enable optimize export model. + """ + from .onnx_export import export + export( + model, + args, + output_path, + input_names=input_names, + output_names=output_names, + opset_version=opset_version, + dynamic_axes=dynamic_axes, + verbose=verbose, + backend=backend, + const_args=const_args, + rewrite_context=rewrite_context, + optimize=optimize) + + @classmethod + def export_from_param(cls, model: Any, param: ONNXParam): + """Export model to ONNX by ONNXParam. + + Examples: + >>> from mmdeploy.ir.onnx import export_from_param + >>> + >>> model = create_model() + >>> param = ONNXParam(...) + >>> + >>> export_from_param(model, param) + + Args: + model (Any): The model to be exported. + param (ONNXParam): The packed export parameter. + """ + from mmdeploy.utils import get_root_logger + logger = get_root_logger() + + # check param validation + param.check() + + # get output path + work_dir = param.work_dir + if not isinstance(work_dir, str): + logger.warning('Invalid work_dir. Use `./work_dir` as default.') + work_dir = './work_dir' + + assert isinstance(param.file_name, str), ('Expect string file name, ' + f'got {type(param.name)}') + output_path = osp.join(param.work_dir, param.file_name) + + cls.export( + model, + param.args, + output_path, + input_names=param.input_names, + output_names=param.output_names, + opset_version=param.opset_version, + dynamic_axes=param.dynamic_axes, + verbose=param.verbose, + backend=param.backend, + const_args=param.const_args, + rewrite_context=param.rewrite_context, + optimize=param.optimize) + + @classmethod + def is_available(cls) -> bool: + """check if the export is available.""" + import importlib + return importlib.util.find_spec('torch') is not None diff --git a/mmdeploy/ir/onnx/onnx_export.py b/mmdeploy/ir/onnx/onnx_export.py new file mode 100644 index 0000000000..cebce72067 --- /dev/null +++ b/mmdeploy/ir/onnx/onnx_export.py @@ -0,0 +1,152 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from copy import deepcopy +from functools import partial +from typing import Any, Dict, List, Optional, Union + +import torch + +from mmdeploy.core import RewriterContext, patch_model +from mmdeploy.utils import get_ir_config, get_root_logger +from mmdeploy.utils.constants import IR, Backend + + +def export(model: Any, + args: Any, + output_path: str, + input_names: Optional[List[str]] = None, + output_names: Optional[List[str]] = None, + opset_version: int = 11, + dynamic_axes: Optional[Dict] = None, + backend: Union[Backend, str] = 'default', + rewrite_context: Optional[Dict] = None, + verbose: bool = False, + const_args: Optional[Dict] = None, + optimize: bool = True): + """export model to ONNX. + + Examples: + >>> from mmdeploy.ir.onnx import export + >>> + >>> model = create_model() + >>> args = get_input_tensor() + >>> + >>> export( + >>> model, + >>> args, + >>> 'place/to/save/model.onnx', + >>> backend='tensorrt', + >>> input_names=['input'], + >>> output_names=['output'], + >>> dynamic_axes={'input': { + >>> 0: 'batch', + >>> 2: 'height', + >>> 3: 'width' + >>> }}) + + Args: + model (Any): Exportable PyTorch Model + args (Any): Arguments are used to trace the graph. + output_path (str): The output path. + input_names (Optional[List[str]], optional): The name of the input in + the graph. Defaults to None. + output_names (Optional[List[str]], optional): The name of the output + in the graph. Defaults to None. + opset_version (int, optional): The ONNX opset version. Defaults to 11. + dynamic_axes (Optional[Dict], optional): Dynamic axes of each inputs. + If not given, all inputs share the fixed shapes of the args. + verbose (bool, optional): Show detail export logs. Defaults to False. + const_args (Optional[Dict], optional): The non-exported inputs of + the model. + rewrite_context (Optional[Dict], optional): The information used by + the rewriter. + optimize (bool): Enable optimize export model. + """ + logger = get_root_logger() + logger.info(f'Export PyTorch model to ONNX: {output_path}.') + + def _add_or_update(cfg: dict, key: str, val: Any): + if key in cfg and isinstance(cfg[key], dict) and isinstance(val, dict): + cfg[key].update(val) + else: + cfg[key] = val + + if rewrite_context is None: + rewrite_context = dict() + + rewrite_context = deepcopy(rewrite_context) + # TODO: deprecate deploy_config format + ir_config = dict( + type='onnx', + input_names=input_names, + output_names=output_names, + opset_version=opset_version, + dynamic_axes=dynamic_axes, + verbose=verbose) + _add_or_update(rewrite_context, 'ir_config', ir_config) + ir = IR.get(get_ir_config(rewrite_context)['type']) + if isinstance(backend, Backend): + backend = backend.value + elif backend is None: + backend = 'default' + backend_config = dict(type=backend) + _add_or_update(rewrite_context, 'backend_config', backend_config) + + # patch model + patched_model = patch_model( + model, cfg=rewrite_context, backend=backend, ir=ir) + + # config optimize info + if backend == Backend.NCNN.value: + optimize = False + if optimize: + from . import optimizer as _optimizer # noqa + if 'onnx_custom_passes' not in rewrite_context: + from .passes import optimize_onnx + onnx_custom_passes = optimize_onnx + else: + onnx_custom_passes = rewrite_context['onnx_custom_passes'] + else: + onnx_custom_passes = None + + # start context + with RewriterContext( + rewrite_context, + backend=backend, + ir=IR.ONNX, + opset=opset_version, + onnx_custom_passes=onnx_custom_passes), torch.no_grad(): + + if const_args is not None: + # patch const_args + assert isinstance( + const_args, dict + ), f'Expect const_args type is dict, get {type(const_args)}.' + model_forward = patched_model.forward + + def wrap_forward(forward): + + def wrapper(*arg, **kwargs): + return forward(*arg, **kwargs) + + return wrapper + + patched_model.forward = wrap_forward(patched_model.forward) + patched_model.forward = partial(patched_model.forward, + **const_args) + + # export with torch.onnx.export + torch.onnx.export( + patched_model, + args, + output_path, + export_params=True, + input_names=input_names, + output_names=output_names, + opset_version=opset_version, + dynamic_axes=dynamic_axes, + keep_initializers_as_inputs=False, + verbose=verbose) + + if const_args is not None: + # recovery forward + patched_model.forward = model_forward diff --git a/mmdeploy/apis/onnx/optimizer.py b/mmdeploy/ir/onnx/optimizer.py similarity index 100% rename from mmdeploy/apis/onnx/optimizer.py rename to mmdeploy/ir/onnx/optimizer.py diff --git a/mmdeploy/apis/onnx/passes/__init__.py b/mmdeploy/ir/onnx/passes/__init__.py similarity index 100% rename from mmdeploy/apis/onnx/passes/__init__.py rename to mmdeploy/ir/onnx/passes/__init__.py diff --git a/mmdeploy/apis/onnx/passes/optimize_onnx.py b/mmdeploy/ir/onnx/passes/optimize_onnx.py similarity index 100% rename from mmdeploy/apis/onnx/passes/optimize_onnx.py rename to mmdeploy/ir/onnx/passes/optimize_onnx.py diff --git a/mmdeploy/ir/torchscript/__init__.py b/mmdeploy/ir/torchscript/__init__.py new file mode 100644 index 0000000000..19c448a01f --- /dev/null +++ b/mmdeploy/ir/torchscript/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .ir_manager import TorchScriptManager, TorchScriptParam + +export = TorchScriptManager.export +export_from_param = TorchScriptManager.export_from_param +is_available = TorchScriptManager.is_available + +__all__ = ['TorchScriptManager', 'TorchScriptParam'] diff --git a/mmdeploy/ir/torchscript/ir_manager.py b/mmdeploy/ir/torchscript/ir_manager.py new file mode 100644 index 0000000000..a0ee5faec5 --- /dev/null +++ b/mmdeploy/ir/torchscript/ir_manager.py @@ -0,0 +1,142 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os.path as osp +from dataclasses import dataclass +from typing import Any, Dict, Optional, Union + +from mmdeploy.utils.constants import Backend +from ..base import IR_MANAGERS, BaseIRManager, BaseIRParam, FileNameDescriptor + + +@dataclass +class TorchScriptParam(BaseIRParam): + """TorchScript IR param. + + Args: + args (Any): The arguments of the model. + work_dir (str): The working directory to save the output. + file_name (str): The file name of the output. postfix can be omitted. + input_names (List[str]): The names to assign to the input of the ir. + output_names (List[str]): The names to assign to the output of the ir. + dynamic_axes (Dict): Determine the dynamic axes of the inputs. It not + given, all axes will be static. + backend (str): The expected backend of the ir. + rewrite_context (Dict): Provide information to the rewriter. + const_args (Any): The constant args of the model. + check_trace (bool): Check outputs after trace. + check_tolerance (float): The tolerance of the check outputs. + """ + # class fields + file_name: FileNameDescriptor = FileNameDescriptor( + default=None, postfix='.pth') + const_args: Any = None + check_trace: bool = True + check_tolerance: float = 1e-05 + + +@IR_MANAGERS.register('torchscript', param=TorchScriptParam) +class TorchScriptManager(BaseIRManager): + """TorchScript IR Manager.""" + + @classmethod + def export(cls, + model: Any, + args: Any, + output_path: str, + backend: Union[Backend, str] = 'default', + rewrite_context: Dict = None, + check_trace: bool = True, + check_tolerance: float = 1e-05, + const_args: Optional[Dict] = None): + """A wrapper of `torch.jit.trace` with some enhancement. + + Examples: + >>> from mmdeploy.ir.torchscript import export + >>> + >>> func = create_model() + >>> inputs = get_input_tensor() + >>> + >>> jit_model = export( + >>> func, + >>> inputs, + >>> backend='torchscript', + >>> check_trace=False) + >>> + + Args: + func (torch.nn.Module): A Python function or `torch.nn.Module` + that will be run with `example_inputs`. + inputs (torch.Tensor, Tuple): A tuple of example inputs that + will be passed to the function while tracing. + output_path (str): The output path. + backend (Backend|str): Which backend will the graph be used. + Different backend would generate different graph. + const_args (Dict): The constant inputs of the model. + rewrite_context (Dict): The information that would be used in + the context of exporting. + check_trace (bool): Check if the same inputs run through traced + code produce the same outputs. + check_tolerance (float): Floating-point comparison tolerance to + use in the checker procedure. + + Returns: + torch.jit.TracedModule: The traced torch jit model. + """ + from .trace import trace + trace( + model, + args, + output_path, + backend=backend, + rewrite_context=rewrite_context, + check_trace=check_trace, + check_tolerance=check_tolerance, + const_args=const_args) + + @classmethod + def export_from_param(cls, model: Any, param: TorchScriptParam): + """Export model to given ir. + + Examples: + >>> from mmdeploy.ir.torchscript import export_from_param + >>> + >>> model = create_model() + >>> param = TorchScriptParam(...) + >>> + >>> export_from_param(model, param) + + Args: + model (Any): The model to be exported. + param (TorchScriptParam): The packed export parameter. + """ + + from mmdeploy.utils import get_root_logger + logger = get_root_logger() + + # check param validation + param.check() + + # get output path + work_dir = param.work_dir + if not isinstance(work_dir, str): + logger.warning('Invalid work_dir. Use `./work_dir` as default.') + work_dir = './work_dir' + + assert isinstance(param.file_name, str), ('Expect string file name, ' + f'got {type(param.name)}') + output_path = osp.join(param.work_dir, param.file_name) + + cls.export( + model, + param.args, + output_path, + backend=param.backend, + rewrite_context=param.rewrite_context, + check_trace=param.check_trace, + check_tolerance=param.check_tolerance, + const_args=param.const_args) + + @classmethod + def is_available(cls) -> bool: + """check if the export is available.""" + import importlib + return importlib.util.find_spec('torch') is not None diff --git a/mmdeploy/ir/torchscript/trace.py b/mmdeploy/ir/torchscript/trace.py new file mode 100644 index 0000000000..c4e5bf4f2c --- /dev/null +++ b/mmdeploy/ir/torchscript/trace.py @@ -0,0 +1,111 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from copy import deepcopy +from functools import partial +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import torch + +from mmdeploy.core import RewriterContext, patch_model +from mmdeploy.utils import IR, Backend, get_ir_config, get_root_logger + + +def trace(func: torch.nn.Module, + inputs: Union[torch.Tensor, Tuple], + output_path: Optional[str] = None, + backend: Union[Backend, str] = 'default', + rewrite_context: Dict = dict(), + check_trace: bool = True, + check_tolerance: float = 1e-05, + const_args: Optional[Dict] = None) -> torch.jit.TracedModule: + """A wrapper of `torch.jit.trace` with some enhancement. + + Examples: + >>> from mmdeploy.ir.torchscript import export + >>> + >>> func = create_model() + >>> inputs = get_input_tensor() + >>> + >>> jit_model = export( + >>> func, + >>> inputs, + >>> output_path, + >>> backend='torchscript', + >>> check_trace=False) + >>> + + Args: + func (torch.nn.Module): A Python function or `torch.nn.Module` that + will be run with `example_inputs`. + inputs (torch.Tensor, Tuple): A tuple of example inputs that will be + passed to the function while tracing. + output_path (str): The output path. + backend (Backend|str): Which backend will the graph be used. Different + backend would generate different graph. + const_args (Dict): The constant inputs of the model. + rewrite_context (Dict): The information that would be used in the + context of exporting. + check_trace (bool): Check if the same inputs run through traced code + produce the same outputs. + check_tolerance (float): Floating-point comparison tolerance to use in + the checker procedure. + + Returns: + torch.jit.TracedModule: The traced torch jit model. + """ + logger = get_root_logger() + logger.info('Export PyTorch model to torchscript.') + + def _add_or_update(cfg: dict, key: str, val: Any): + if key in cfg and isinstance(cfg[key], dict) and isinstance(val, dict): + cfg[key].update(val) + else: + cfg[key] = val + + if rewrite_context is None: + rewrite_context = dict() + + rewrite_context = deepcopy(rewrite_context) + ir_config = dict(type='torchscript') + _add_or_update(rewrite_context, 'ir_config', ir_config) + + if isinstance(backend, Backend): + backend = backend.value + elif backend is None: + backend = 'default' + backend_config = dict(type=backend) + _add_or_update(rewrite_context, 'backend_config', backend_config) + + # patch model + if isinstance(func, torch.nn.Module): + ir = IR.get(get_ir_config(rewrite_context)['type']) + func = patch_model(func, cfg=rewrite_context, backend=backend, ir=ir) + + with RewriterContext( + rewrite_context, ir=IR.TORCHSCRIPT, + backend=backend), torch.no_grad(): + + # patch const_args + if const_args is not None: + assert isinstance( + const_args, dict + ), f'Expect const_args type is dict, get {type(const_args)}.' + model_forward = func.forward + func.forward = partial(func.forward, **const_args) + + # for exporting models with weight that depends on inputs + func(*inputs) if isinstance(inputs, Sequence) \ + else func(inputs) + ts_model = torch.jit.trace( + func, + inputs, + check_trace=check_trace, + check_tolerance=check_tolerance) + + if const_args is not None: + func.forward = model_forward + + # save model + logger.info(f'Save PyTorch model: {output_path}.') + torch.jit.save(ts_model, output_path) + + return ts_model diff --git a/mmdeploy/tools/__init__.py b/mmdeploy/tools/__init__.py new file mode 100644 index 0000000000..ef101fec61 --- /dev/null +++ b/mmdeploy/tools/__init__.py @@ -0,0 +1 @@ +# Copyright (c) OpenMMLab. All rights reserved. diff --git a/mmdeploy/tools/console.py b/mmdeploy/tools/console.py new file mode 100644 index 0000000000..66200bb452 --- /dev/null +++ b/mmdeploy/tools/console.py @@ -0,0 +1,197 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import importlib +import sys +from argparse import ArgumentParser +from typing import Iterable, List, Optional + + +def import_custom_modules(custom_modules: Iterable): + """Import custom module.""" + from mmdeploy.utils import get_root_logger + logger = get_root_logger(0) + custom_modules = [] if custom_modules is None else custom_modules + + for qualname in custom_modules: + try: + importlib.import_module(qualname) + logger.info(f'Import custom module: {qualname}') + except Exception as e: + logger.warning('Failed to import custom module: ' + f'{qualname} with error: {e}') + + +def list_command(parser: ArgumentParser, + input_args: Optional[List[str]] = None): + """List command.""" + parser.description = 'List available backend and task.' + arg_group = parser.add_argument_group('List Options') + arg_group.add_argument( + '--backend', action='store_true', help='List available backend.') + arg_group.add_argument( + '--task', action='store_true', help='List available task.') + arg_group.add_argument( + '--custom-modules', type=str, nargs='*', help='Custom module path.') + args = parser.parse_args(input_args) + + def _print_pretty_table(data_list, title): + """print pretty table.""" + max_data_length = [len(t) for t in title] + + for data in data_list: + assert len(title) == len(data) + for idx, d in enumerate(data): + max_data_length[idx] = max(max_data_length[idx], len(d)) + + # print title + title_all = ' '.join([ + f'{t}' + (' ' * (length - len(t))) + for length, t in zip(max_data_length, title) + ]) + print(title_all) + + # print dash + dash_all = ' '.join(['-' * length for length in max_data_length]) + print(dash_all) + + # print data + for data in data_list: + data_all = ' '.join([ + f'{t}' + (' ' * (length - len(t))) + for length, t in zip(max_data_length, data) + ]) + print(data_all) + + # custom import + import_custom_modules(args.custom_modules) + enable_list_backend = args.backend + enable_list_task = args.task + + if not enable_list_backend and not enable_list_task: + enable_list_task = True + enable_list_backend = True + + if enable_list_backend: + from mmdeploy.backend.base import get_backend_manager + from mmdeploy.utils import Backend + + exclude_backend_lists = [Backend.DEFAULT, Backend.PYTORCH, Backend.SDK] + backend_lists = [ + backend.value for backend in Backend + if backend not in exclude_backend_lists + ] + + # get all available backend + available_backend = [] + for backend in backend_lists: + backend_mgr = get_backend_manager(backend) + if backend_mgr.is_available(): + try: + available_backend.append( + (backend, backend_mgr.get_version())) + except Exception as e: + sys.stderr.write(f'List backend: {backend} failed.' + f' with error: {e}\n') + + _print_pretty_table(available_backend, ['Backend', 'Version']) + + if enable_list_task: + # TODO: add list task here + pass + + +def show_command(parser: ArgumentParser, + input_args: Optional[List[str]] = None): + """Show command.""" + parser.description = 'Show environment of the backend or task.' + arg_group = parser.add_argument_group('Show Options') + arg_group.add_argument('name', help='The object name to show.') + arg_group.add_argument( + '--custom-modules', type=str, nargs='*', help='Custom module path.') + args = parser.parse_args(input_args) + + # custom import + import_custom_modules(args.custom_modules) + + obj_name = args.name + + # Check if obj is backend + from mmdeploy.utils import Backend + + exclude_backend_lists = [Backend.DEFAULT, Backend.PYTORCH, Backend.SDK] + backend_lists = [ + backend.value for backend in Backend + if backend not in exclude_backend_lists + ] + + if obj_name in backend_lists: + from mmdeploy.backend.base import get_backend_manager + backend_mgr = get_backend_manager(obj_name) + if backend_mgr.is_available(): + backend_mgr.check_env(print) + else: + sys.stderr.write(f'Backend: {obj_name} is not available.\n') + + # TODO: add show task here + + +def run_command(parser: ArgumentParser, + input_args: Optional[List[str]] = None): + """Run command.""" + + # extract help + help = False + if '-h' in input_args: + help = True + input_args.remove('-h') + if '--help' in input_args: + help = True + input_args.remove('--help') + + parser.description = 'Run console tools of backend or task.' + arg_group = parser.add_argument_group('Run Options') + arg_group.add_argument( + 'obj_name', + default=None, + help='The backend or task name to run the console tools.') + arg_group.add_argument( + '--custom-modules', type=str, nargs='*', help='Custom module path.') + + if help: + if len(input_args) == 0 or input_args[0] == '--custom-modules': + parser.print_help() + exit() + + args, remain_args = parser.parse_known_args(input_args) + + # custom import + import_custom_modules(args.custom_modules) + + obj_name = args.obj_name + + if help: + remain_args += ['--help'] + + # Check if obj is backend + from mmdeploy.utils import Backend + + exclude_backend_lists = [Backend.DEFAULT, Backend.PYTORCH, Backend.SDK] + backend_lists = [ + backend.value for backend in Backend + if backend not in exclude_backend_lists + ] + + if obj_name in backend_lists: + from mmdeploy.backend.base import get_backend_manager + backend_mgr = get_backend_manager(obj_name) + if backend_mgr.is_available(): + parser = ArgumentParser() + try: + with backend_mgr.parse_args(parser, remain_args): + pass + except NotImplementedError: + sys.stderr.write(f'Backend: {obj_name} console tools' + ' has not been implemented.\n') + else: + sys.stderr.write(f'Backend: {obj_name} is not available.\n') + + # TODO: add run task here diff --git a/mmdeploy/utils/docstring_parser.py b/mmdeploy/utils/docstring_parser.py new file mode 100644 index 0000000000..91a95e39f8 --- /dev/null +++ b/mmdeploy/utils/docstring_parser.py @@ -0,0 +1,253 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import inspect +import re +from dataclasses import dataclass +from typing import Any, List, Optional, Tuple + + +@dataclass +class DocStrArg: + """The argument of the docstring.""" + name: str + type: str + desc: str + + +@dataclass +class DocStr: + """The packed docstring description.""" + head: str + desc: str + args: List[DocStrArg] + + +def parse_empty_line(buffer: str) -> int: + """Get empty line in the buffer. + + Args: + buffer (str): The string buffer. + + Returns: + int: The length of the empty lines. + """ + pattern = r'^[^\S\n]*(\n|$)' + + ret_len = 0 + while len(buffer) > 0: + m = re.match(pattern, buffer) + if m is None: + break + next_len = m.end() - m.start() + ret_len += next_len + buffer = buffer[next_len:] + + return ret_len + + +def parse_arg(buffer: str) -> Tuple[Optional[DocStrArg], int]: + """Parse one argument at the beginning of the buffer. + + Args: + buffer (str): The docstring buffer + + Returns: + Tuple[Optional[DocStrArg], int]: The argument info and the + buffer length. If there is no argument, return (None, 0) + """ + ret_len = 0 + pattern = r'^[^\S\n]*(?P\w+)[^\S\n]*\((?P.*)\):[^\S\n]*(?P.*)(\n|$)' # noqa + m = re.match(pattern, buffer) + if m is None: + return None, 0 + + doc_dict = m.groupdict() + next_len = m.end() - m.start() + buffer = buffer[next_len:] + ret_len += next_len + + # try read docs with multiline + while len(buffer) > 0: + # if next line is starts with same pattern, it is not remain + m = re.match(pattern, buffer) + if m is not None: + break + + # remain should starts with whitespace. + m = re.match(r'^[^\S\n]+(?P\S.*)(\n|$)', buffer) + if m is None: + break + remain_desc = m.group('desc') + doc_dict['desc'] += remain_desc + next_len = m.end() - m.start() + buffer = buffer[next_len:] + ret_len += next_len + + doc_arg = DocStrArg(**doc_dict) + return doc_arg, ret_len + + +def parse_args(buffer: str) -> Tuple[List[DocStrArg], int]: + """Parse all arguments at the beginning of the buffer. + + Args: + buffer (str): The docstring buffer + + Returns: + Tuple[List[DocStrArg], int]: The list of all parsed arguments and + buffer length. + """ + ret_args = [] + ret_len = 0 + while len(buffer) > 0: + doc_arg, doc_len = parse_arg(buffer) + if doc_len == 0 or doc_arg is None: + break + ret_args.append(doc_arg) + buffer = buffer[doc_len:] + ret_len += doc_len + + return ret_args, ret_len + + +def parse_args_section(buffer: str) -> Tuple[List[DocStrArg], int]: + """Parse The arguments section, with the head `Args:` + + Args: + buffer (str): The docstring buffer + + Returns: + Tuple[List[DocStrArg], int]: The list of all parsed arguments and + buffer length. + """ + ret_len = 0 + + def _update_buffer(next_len: int): + nonlocal buffer + nonlocal ret_len + buffer = buffer[next_len:] + ret_len += next_len + + def _skip_empty_line(): + nonlocal buffer + next_len = parse_empty_line(buffer) + _update_buffer(next_len) + + # parse `Args:` + head_pattern = r'^Args:\s*(\n|$)' + m = re.match(head_pattern, buffer) + if m is None: + return [], 0 + next_len = m.end() - m.start() + _update_buffer(next_len) + _skip_empty_line() + + # parse args + doc_args, next_len = parse_args(buffer) + _update_buffer(next_len) + _skip_empty_line() + + return doc_args, ret_len + + +SECTION_HEAD = ['Args:', 'Returns:', 'Examples:'] + + +def parse_desc(buffer: str) -> Tuple[str, int]: + """Parse the description of the docstring. + + Args: + buffer (str): The docstring buffer + + Returns: + Tuple[str, int]: The description string and the buffer length + """ + desc_pattern = r'^(?P.*(\n|$))' + desc = '' + desc_len = 0 + while len(buffer) > 0: + m = re.match(desc_pattern, buffer) + if m is None: + break + line_desc = m.group('desc') + + # check if buffer reach next section + if line_desc.rstrip() in SECTION_HEAD: + break + + line_len = m.end() - m.start() + desc += line_desc + desc_len += line_len + buffer = buffer[line_len:] + + return desc.strip(), desc_len + + +def parse_docstring(buffer: str) -> Optional[DocStr]: + """Parse docstring. + + Args: + buffer (str): The docstring buffer + + Returns: + Optional[DocStr]: The parsed docstring info. If parse failed. return + None. + """ + head = '' + desc = '' + args = [] + + def _skip_empty_line(): + nonlocal buffer + next_len = parse_empty_line(buffer) + buffer = buffer[next_len:] + + # parse head + head_pattern = r'^(?P.*)(\n|$)' + m = re.match(head_pattern, buffer) + if m is None: + return DocStr(head=head, desc=desc, args=args) + next_len = m.end() - m.start() + head = m.group('head').rstrip() + buffer = buffer[next_len:] + _skip_empty_line() + + # parse desc + desc, next_len = parse_desc(buffer) + buffer = buffer[next_len:] + + while len(buffer) > 0: + section_pattern = r'^(?P
.*)(\n|$)' + m = re.match(section_pattern, buffer) + if m is None: + return DocStr(head=head, desc=desc, args=args) + + if m.group('section').rstrip() == 'Args:': + args, next_len = parse_args_section(buffer) + else: + next_len = m.end() - m.start() + + buffer = buffer[next_len:] + + return DocStr(head=head, desc=desc, args=args) + + +def inspect_docstring_arguments(obj: Any, + ignore_args: Optional[List[str]] = None + ) -> List[DocStrArg]: + """Inspect google style function docstring. + + Args: + obj (Any): Object with docstring + ignore_args (List[str]): Ignore arguments when parsing. + + Returns: + List[DocStrArg]: Parsed docstring arguments. + """ + ignore_args = [] if ignore_args is None else ignore_args + doc_str: DocStr = parse_docstring(inspect.getdoc(obj)) + + def _filter_cb(arg: DocStrArg): + nonlocal ignore_args + return arg.name not in ignore_args + + return list(filter(_filter_cb, doc_str.args)) diff --git a/mmdeploy/utils/test.py b/mmdeploy/utils/test.py index 835f5b8b57..42863de9eb 100644 --- a/mmdeploy/utils/test.py +++ b/mmdeploy/utils/test.py @@ -1,6 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. - -import os.path as osp import random import string import tempfile @@ -147,7 +145,7 @@ class SwitchBackendWrapper: """A switcher for backend wrapper for unit tests. Examples: >>> from mmdeploy.utils.test import SwitchBackendWrapper - >>> from mmdeploy.backend.onnxruntime import ORTWrapper + >>> from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper >>> with SwitchBackendWrapper(ORTWrapper) as wrapper: >>> wrapper.set(ORTWrapper, outputs=outputs) >>> ... @@ -184,8 +182,10 @@ def __call__(self, *args, **kwds): def __init__(self, recover_class): self._recover_class = recover_class + self._initialized = False def __enter__(self): + self.set() return self def __exit__(self, type, value, trace): @@ -194,12 +194,14 @@ def __exit__(self, type, value, trace): def set(self, **kwargs): """Replace attributes in backend wrappers with dummy items.""" obj = self._recover_class - self.init = obj.__init__ - self.forward = obj.forward - self.call = obj.__call__ - obj.__init__ = SwitchBackendWrapper.BackendWrapper.__init__ - obj.forward = SwitchBackendWrapper.BackendWrapper.forward - obj.__call__ = SwitchBackendWrapper.BackendWrapper.__call__ + if not self._initialized: + self.init = obj.__init__ + self.forward = obj.forward + self.call = obj.__call__ + obj.__init__ = SwitchBackendWrapper.BackendWrapper.__init__ + obj.forward = SwitchBackendWrapper.BackendWrapper.forward + obj.__call__ = SwitchBackendWrapper.BackendWrapper.__call__ + self._initialized = True for k, v in kwargs.items(): setattr(obj, k, v) @@ -208,10 +210,12 @@ def recover(self): assert self.init is not None and \ self.forward is not None,\ 'recover method must be called after exchange' - obj = self._recover_class - obj.__init__ = self.init - obj.forward = self.forward - obj.__call__ = self.call + if self._initialized: + obj = self._recover_class + obj.__init__ = self.init + obj.forward = self.forward + obj.__call__ = self.call + self._initialized = False def assert_allclose(expected: List[Union[torch.Tensor, np.ndarray]], @@ -333,17 +337,17 @@ def forward(self, inputs: dict): model = DummyModel().eval() - with RewriterContext( - cfg=deploy_cfg, backend=backend.value, opset=11), torch.no_grad(): - torch.onnx.export( + with torch.no_grad(): + from mmdeploy.ir.onnx import ONNXManager + ONNXManager.export( model, (model_inputs, {}), onnx_file_path, - export_params=True, input_names=input_names, output_names=output_names, opset_version=11, dynamic_axes=dynamic_axes, - keep_initializers_as_inputs=False) + backend=backend.value, + rewrite_context=deploy_cfg) return onnx_file_path @@ -360,20 +364,18 @@ def get_ts_model(wrapped_model: nn.Module, Returns: str: The path to the TorchScript model file. """ - ir_file_path = tempfile.NamedTemporaryFile(suffix='.pt').name + ir_file_path = tempfile.NamedTemporaryFile(suffix='.pth').name backend = get_backend(deploy_cfg) - from mmdeploy.apis.torch_jit import trace - context_info = dict(deploy_cfg=deploy_cfg) - output_prefix = osp.splitext(ir_file_path)[0] + from mmdeploy.ir.torchscript import TorchScriptManager example_inputs = [v for _, v in model_inputs.items()] - trace( + TorchScriptManager.export( wrapped_model, example_inputs, - output_path_prefix=output_prefix, + output_path=ir_file_path, backend=backend, - context_info=context_info) + rewrite_context=deploy_cfg) return ir_file_path @@ -404,44 +406,45 @@ def get_backend_outputs(ir_file_path: str, k for k, v in flatten_model_inputs.items() if k != 'ctx' ] - work_dir = tempfile.TemporaryDirectory().name - device = 'cpu' - - # TODO: Try to wrap these code into backend manager - if backend != Backend.TORCHSCRIPT: - model_inputs = flatten_model_inputs - if backend == Backend.TENSORRT: - device = 'cuda' - model_inputs = dict((k, v.cuda()) for k, v in model_inputs.items()) - elif backend == Backend.OPENVINO: - input_info = { - name: value.shape - for name, value in flatten_model_inputs.items() - } - deploy_cfg['backend_config']['model_inputs'] = [ - dict(opt_shapes=input_info) - ] - backend_files = to_backend( - backend.value, [ir_file_path], - work_dir=work_dir, - deploy_cfg=deploy_cfg, - device=device) - backend_feats = model_inputs - - if backend == Backend.TORCHSCRIPT: - backend_feats = [v for _, v in model_inputs.items()] - - from mmdeploy.codebase.base import BaseBackendModel - backend_model = BaseBackendModel._build_wrapper( - backend, - backend_files, - device, - input_names=input_names, - output_names=output_names) - with torch.no_grad(): - backend_outputs = backend_model(backend_feats) - backend_outputs = backend_model.output_to_list(backend_outputs) - return backend_outputs + with tempfile.TemporaryDirectory() as work_dir: + device = 'cpu' + + # TODO: Try to wrap these code into backend manager + if backend != Backend.TORCHSCRIPT: + model_inputs = flatten_model_inputs + if backend == Backend.TENSORRT: + device = 'cuda' + model_inputs = dict((k, v.cuda()) for k, v in model_inputs.items()) + elif backend == Backend.OPENVINO: + input_info = { + name: value.shape + for name, value in flatten_model_inputs.items() + } + deploy_cfg['backend_config']['model_inputs'] = [ + dict(opt_shapes=input_info) + ] + backend_files = to_backend( + backend.value, [ir_file_path], + work_dir=work_dir, + deploy_cfg=deploy_cfg, + device=device) + backend_feats = model_inputs + + if backend == Backend.TORCHSCRIPT: + backend_feats = [v for _, v in model_inputs.items()] + + from mmdeploy.codebase.base import BaseBackendModel + backend_model = BaseBackendModel._build_wrapper( + backend, + backend_files, + device, + input_names=input_names, + output_names=output_names, + deploy_cfg=deploy_cfg) + with torch.no_grad(): + backend_outputs = backend_model(backend_feats) + backend_outputs = backend_model.output_to_list(backend_outputs) + return backend_outputs def get_rewrite_outputs(wrapped_model: nn.Module, diff --git a/tests/test_apis/test_onnx2ascend.py b/tests/test_apis/test_onnx2ascend.py index fd853936f2..1d54e3c2bc 100644 --- a/tests/test_apis/test_onnx2ascend.py +++ b/tests/test_apis/test_onnx2ascend.py @@ -66,6 +66,5 @@ def test_onnx2ascend(): dict( dynamic_batch_size=[1, 2, 4], input_shapes=dict(input=[-1, 3, 224, 224]))) - from_onnx(onnx_file, work_dir, model_inputs) - assert osp.exists(work_dir) + from_onnx(onnx_file, file_name, model_inputs) assert osp.exists(om_path) diff --git a/tests/test_apis/test_onnx2ncnn.py b/tests/test_apis/test_onnx2ncnn.py index ef6cf1f1de..f9f3248467 100644 --- a/tests/test_apis/test_onnx2ncnn.py +++ b/tests/test_apis/test_onnx2ncnn.py @@ -61,8 +61,7 @@ def test_onnx2ncnn(): work_dir, _ = osp.split(onnx_file) save_param, save_bin = get_output_model_file(onnx_file, work_dir=work_dir) - file_name = osp.splitext(onnx_file)[0] - from_onnx(onnx_file, osp.join(work_dir, file_name)) + from_onnx(onnx_file, save_param, save_bin) assert osp.exists(work_dir) assert osp.exists(save_param) assert osp.exists(save_bin) diff --git a/tests/test_apis/test_onnx2openvino.py b/tests/test_apis/test_onnx2openvino.py index 0a9302f16a..6a8d32fbb2 100644 --- a/tests/test_apis/test_onnx2openvino.py +++ b/tests/test_apis/test_onnx2openvino.py @@ -53,7 +53,7 @@ def get_outputs(pytorch_model, openvino_model_path, input, input_name, output_name): output_pytorch = pytorch_model(input).numpy() - from mmdeploy.backend.openvino import OpenVINOWrapper + from mmdeploy.backend.openvino.wrapper import OpenVINOWrapper openvino_model = OpenVINOWrapper(openvino_model_path) openvino_output = openvino_model({input_name: input})[output_name] @@ -92,21 +92,29 @@ def test_onnx2openvino(get_deploy_cfg): input_info = {input_name: export_img.shape} output_names = [output_name] - openvino_dir = tempfile.TemporaryDirectory().name - deploy_cfg = get_deploy_cfg() - mo_options = get_mo_options_from_cfg(deploy_cfg) - from_onnx(onnx_file, openvino_dir, input_info, output_names, 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.' - - test_img = torch.rand([1, 3, 16, 16]) - output_pytorch, openvino_output = get_outputs(pytorch_model, - openvino_model_path, - test_img, input_name, - output_name) - assert np.allclose(output_pytorch, openvino_output), \ - 'OpenVINO and PyTorch outputs are not the same.' + with tempfile.TemporaryDirectory() as openvino_dir: + deploy_cfg = get_deploy_cfg() + mo_options = get_mo_options_from_cfg(deploy_cfg) + mo_options = mo_options.get_options() + openvino_model_path = get_output_model_file(onnx_file, openvino_dir) + + from_onnx( + onnx_file, + openvino_model_path, + input_info, + output_names, + work_dir=openvino_dir, + mo_options=mo_options) + assert osp.exists(openvino_model_path), \ + 'The file (.xml) for OpenVINO IR has not been created.' + + test_img = torch.rand([1, 3, 16, 16]) + output_pytorch, openvino_output = get_outputs(pytorch_model, + openvino_model_path, + test_img, input_name, + output_name) + assert np.allclose(output_pytorch, openvino_output), \ + 'OpenVINO and PyTorch outputs are not the same.' @backend_checker(Backend.OPENVINO) diff --git a/tests/test_apis/test_onnx2tensorrt.py b/tests/test_apis/test_onnx2tensorrt.py index c84399530f..9117ac47c9 100644 --- a/tests/test_apis/test_onnx2tensorrt.py +++ b/tests/test_apis/test_onnx2tensorrt.py @@ -75,7 +75,7 @@ def generate_onnx_file(model): @backend_checker(Backend.TENSORRT) def test_onnx2tensorrt(): from mmdeploy.apis.tensorrt import onnx2tensorrt - from mmdeploy.backend.tensorrt import load + from mmdeploy.backend.tensorrt.utils import load model = test_model generate_onnx_file(model) deploy_cfg = get_deploy_cfg() diff --git a/tests/test_apis/test_onnx_passes.py b/tests/test_apis/test_onnx_passes.py index d3b1a83e23..3a8b99235a 100644 --- a/tests/test_apis/test_onnx_passes.py +++ b/tests/test_apis/test_onnx_passes.py @@ -7,9 +7,8 @@ import torch import torch.nn as nn -from mmdeploy.apis.onnx.optimizer import \ - model_to_graph__custom_optimizer # noqa from mmdeploy.core import RewriterContext +from mmdeploy.ir.onnx.optimizer import model_to_graph__custom_optimizer # noqa onnx_file = tempfile.NamedTemporaryFile(suffix='.onnx').name @@ -189,9 +188,6 @@ def forward(self, x): node, idx = _find_next_node(0, nodes, 'GlobalAveragePool') assert node is not None - node, idx = _find_next_node(idx + 1, nodes, 'Flatten') - assert node is not None - def test_fuse_select_assign(): pytest.importorskip('mmdeploy.backend.torchscript.ts_optimizer.onnx') diff --git a/tests/test_apis/test_torch2torchscript.py b/tests/test_apis/test_torch2torchscript.py index 364f1be0f3..fbb067e86b 100644 --- a/tests/test_apis/test_torch2torchscript.py +++ b/tests/test_apis/test_torch2torchscript.py @@ -10,7 +10,7 @@ from mmdeploy.utils import IR, Backend from mmdeploy.utils.test import get_random_name -ts_file = tempfile.NamedTemporaryFile(suffix='.pt').name +ts_file = tempfile.NamedTemporaryFile(suffix='.pth').name input_name = get_random_name() output_name = get_random_name() diff --git a/tests/test_backend/conftest.py b/tests/test_backend/conftest.py new file mode 100644 index 0000000000..111f8705f6 --- /dev/null +++ b/tests/test_backend/conftest.py @@ -0,0 +1,208 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from tempfile import NamedTemporaryFile + +import pytest + + +@pytest.fixture(scope='module') +def input_shape(): + yield [1, 3, 8, 8] + + +@pytest.fixture(scope='module') +def dummy_x(input_shape): + torch = pytest.importorskip('torch') + yield torch.rand(input_shape) + + +@pytest.fixture(scope='module') +def dummy_y(input_shape): + torch = pytest.importorskip('torch') + yield torch.rand(input_shape) + + +# ------ 1i1o ------ + + +@pytest.fixture(scope='module') +def dummy_torch_model_1i1o(): + torch = pytest.importorskip('torch') + + class DummyModel(torch.nn.Module): + + def forward(self, x): + return x + x + + yield DummyModel() + + +@pytest.fixture(scope='module') +def input_1i(dummy_x): + yield dummy_x + + +@pytest.fixture(scope='module') +def output_1i1o(dummy_torch_model_1i1o, input_1i): + yield dummy_torch_model_1i1o(input_1i) + + +@pytest.fixture(scope='module') +def input_names_1i(): + yield ['x'] + + +@pytest.fixture(scope='module') +def output_names_1i1o(): + yield ['ox'] + + +@pytest.fixture(scope='module') +def input_dict_1i(input_1i, input_names_1i): + yield dict(zip(input_names_1i, [input_1i])) + + +@pytest.fixture(scope='module') +def output_dict_1i1o(output_1i1o, output_names_1i1o): + yield dict(zip(output_names_1i1o, [output_1i1o])) + + +@pytest.fixture(scope='module') +def input_shape_dict_1i(input_dict_1i): + yield dict((name, val.shape) for name, val in input_dict_1i.items()) + + +@pytest.fixture(scope='module') +def dynamic_axes_1i(): + yield dict(x={0: 'b', 2: 'h', 3: 'w'}) + + +@pytest.fixture(scope='module') +def onnx_model_static_1i1o(dummy_torch_model_1i1o, input_1i, input_names_1i, + output_names_1i1o): + torch = pytest.importorskip('torch') + tmp_path = NamedTemporaryFile(suffix='.onnx').name + torch.onnx.export( + dummy_torch_model_1i1o, + input_1i, + tmp_path, + input_names=input_names_1i, + output_names=output_names_1i1o) + + yield tmp_path + + +# ------ 2i2o ------ + + +@pytest.fixture(scope='module') +def dummy_torch_model_2i2o(): + torch = pytest.importorskip('torch') + + class DummyModel(torch.nn.Module): + + def forward(self, x, y): + return x + y, x * y + + yield DummyModel() + + +@pytest.fixture(scope='module') +def input_2i(dummy_x, dummy_y): + yield dummy_x, dummy_y + + +@pytest.fixture(scope='module') +def output_2i2o(dummy_torch_model_2i2o, dummy_x, dummy_y): + yield dummy_torch_model_2i2o(dummy_x, dummy_y) + + +@pytest.fixture(scope='module') +def input_names_2i(): + yield ['x', 'y'] + + +@pytest.fixture(scope='module') +def output_names_2i2o(): + yield ['ox', 'oy'] + + +@pytest.fixture(scope='module') +def input_dict_2i(input_2i, input_names_2i): + yield dict(zip(input_names_2i, input_2i)) + + +@pytest.fixture(scope='module') +def output_dict_2i2o(output_2i2o, output_names_2i2o): + yield dict(zip(output_names_2i2o, output_2i2o)) + + +@pytest.fixture(scope='module') +def input_shape_dict_2i(input_dict_2i): + yield dict((name, val.shape) for name, val in input_dict_2i.items()) + + +@pytest.fixture(scope='module') +def output_shape_dict_2i2o(output_dict_2i2o): + yield dict((name, val.shape) for name, val in output_dict_2i2o.items()) + + +@pytest.fixture(scope='module') +def dynamic_axes_2i(): + yield dict(x={0: 'b', 2: 'h', 3: 'w'}, y={0: 'b', 2: 'h', 3: 'w'}) + + +@pytest.fixture(scope='module') +def onnx_model_static_2i2o(dummy_torch_model_2i2o, dummy_x, dummy_y, + input_names_2i, output_names_2i2o): + torch = pytest.importorskip('torch') + tmp_path = NamedTemporaryFile(suffix='.onnx').name + torch.onnx.export( + dummy_torch_model_2i2o, (dummy_x, dummy_y), + tmp_path, + input_names=input_names_2i, + output_names=output_names_2i2o) + + yield tmp_path + + +@pytest.fixture(scope='module') +def onnx_model_dynamic_2i2o(dummy_torch_model_2i2o, dummy_x, dummy_y, + input_names_2i, output_names_2i2o, + dynamic_axes_2i): + torch = pytest.importorskip('torch') + tmp_path = NamedTemporaryFile(suffix='.onnx').name + torch.onnx.export( + dummy_torch_model_2i2o, (dummy_x, dummy_y), + tmp_path, + input_names=input_names_2i, + output_names=output_names_2i2o, + dynamic_axes=dynamic_axes_2i) + + yield tmp_path + + +@pytest.fixture(scope='module') +def torchscript_model2i2o(dummy_torch_model_2i2o, dummy_x, dummy_y): + torch = pytest.importorskip('torch') + tmp_path = NamedTemporaryFile(suffix='.pth').name + jit_model = torch.jit.trace(dummy_torch_model_2i2o, (dummy_x, dummy_y)) + torch.jit.save(jit_model, tmp_path) + + yield tmp_path + + +@pytest.fixture +def assert_forward(): + try: + from torch.testing import assert_close as torch_assert_close + except Exception: + from torch.testing import assert_allclose as torch_assert_close + + def _impl(model, inputs, gts, rtol=None, atol=None): + outputs = model(inputs) + for name in outputs: + out = outputs[name] + gt = gts[name] + torch_assert_close(out, gt, rtol=rtol, atol=atol) + + return _impl diff --git a/tests/test_backend/test_ascend.py b/tests/test_backend/test_ascend.py new file mode 100644 index 0000000000..e83c41df42 --- /dev/null +++ b/tests/test_backend/test_ascend.py @@ -0,0 +1,127 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.ascend import AscendManager as backend_mgr +from mmdeploy.backend.ascend import AscendParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.om' + + +class TestBackendParam: + + def test_get_model_files(self): + param = AscendParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + + assert param.get_model_files() == 'tmp' + _extension + + def test_add_argument(self): + parser = argparse.ArgumentParser() + AscendParam.add_argument( + parser, 'dynamic_dims', dtype=str, default=None, desc='') + + assert parser.parse_args(['--dynamic-dims', + '3,2,-1']).dynamic_dims == { + None: (3, 2, -1) + } + assert parser.parse_args(['--dynamic-dims', + 'input:3,2,-1']).dynamic_dims == { + 'input': (3, 2, -1) + } + assert parser.parse_args( + ['--dynamic-dims', 'input1:3,2,-1', + 'input2:3,2,-1']).dynamic_dims == { + 'input1': (3, 2, -1), + 'input2': (3, 2, -1) + } + assert parser.parse_args( + ['--dynamic-dims', + 'input1:3,2,-1;input2:3,2,-1']).dynamic_dims == { + 'input1': (3, 2, -1), + 'input2': (3, 2, -1) + } + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_2i): + yield input_shape_dict_2i + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_static_2i2o): + yield onnx_model_static_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model, input_shape_dict): + from mmdeploy.backend.ascend import AtcParam + with TemporaryDirectory() as tmp_dir: + model_path = osp.join(tmp_dir, 'tmp' + _extension) + backend_mgr.to_backend(onnx_model, model_path, + AtcParam(input_shapes=input_shape_dict)) + + yield model_path + + def test_to_backend(self, backend_model): + assert osp.exists(backend_model) + + def test_to_backend_from_param(self, onnx_model, input_shape_dict): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_shapes=input_shape_dict) + backend_mgr.to_backend_from_param(onnx_model, param) + + param_path = param.get_model_files() + assert osp.exists(param_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, + assert_forward): + wrapper = backend_mgr.build_wrapper(backend_model) + assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + assert_forward): + param = backend_mgr.build_param( + work_dir='', device='cpu', file_name=backend_model) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + def test_parse_args(self, onnx_model, input_shape_dict): + + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + _extension + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + args += ['--input-shapes', input_shapes] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, param_name)) diff --git a/tests/test_backend/test_coreml.py b/tests/test_backend/test_coreml.py new file mode 100644 index 0000000000..2f14c20469 --- /dev/null +++ b/tests/test_backend/test_coreml.py @@ -0,0 +1,118 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.coreml import CoreMLManager as backend_mgr +from mmdeploy.backend.coreml import CoreMLParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.mlpackage' + + +class TestBackendParam: + + def test_get_model_files(self): + param = CoreMLParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + + assert param.get_model_files() == 'tmp' + _extension + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_2i): + yield input_shape_dict_2i + + @pytest.fixture(scope='class') + def input_names(self, input_names_2i): + yield input_names_2i + + @pytest.fixture(scope='class') + def output_names(self, output_names_2i2o): + yield output_names_2i2o + + @pytest.fixture(scope='class') + def ir_model(self, torchscript_model2i2o): + yield torchscript_model2i2o + + @pytest.fixture(scope='class') + def backend_model(self, ir_model, input_names, output_names, + input_shape_dict): + with TemporaryDirectory() as tmp_dir: + model_path = osp.join(tmp_dir, 'tmp' + _extension) + backend_mgr.to_backend( + ir_model, + model_path, + input_names=input_names, + output_names=output_names, + input_shapes=input_shape_dict) + + yield model_path + + def test_to_backend(self, backend_model): + assert osp.exists(backend_model) + + def test_to_backend_from_param(self, ir_model, input_names, output_names, + input_shape_dict): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_shapes=input_shape_dict, + input_names=input_names, + output_names=output_names) + backend_mgr.to_backend_from_param(ir_model, param) + + param_path = param.get_model_files() + assert osp.exists(param_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, + assert_forward): + wrapper = backend_mgr.build_wrapper(backend_model) + assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + assert_forward): + param = backend_mgr.build_param(work_dir='', file_name=backend_model) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + def test_parse_args(self, ir_model, input_names, output_names, + input_shape_dict): + + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + _extension + # make args + args = ['convert'] + args += ['--torchscript-path', ir_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + args += ['--input-shapes', input_shapes] + args += ['--input-names', *input_names] + args += ['--output-names', *output_names] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, param_name)) diff --git a/tests/test_backend/test_manager.py b/tests/test_backend/test_manager.py new file mode 100644 index 0000000000..49111116a2 --- /dev/null +++ b/tests/test_backend/test_manager.py @@ -0,0 +1,180 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse + +import pytest + +from mmdeploy.backend.base import BaseBackendParam + + +def test_parse_shape(): + + parser = argparse.ArgumentParser() + BaseBackendParam.add_arguments(parser) + + args = parser.parse_args(['--input-shapes', 'input:1x2x3x4']) + assert args.input_shapes == dict(input=[1, 2, 3, 4]) + + args = parser.parse_args( + ['--input-shapes', 'input1:1x2x3x4,input2:4x3x2x1']) + assert args.input_shapes == dict(input1=[1, 2, 3, 4], input2=[4, 3, 2, 1]) + + args = parser.parse_args( + ['--input-shapes', 'input1:1x2x3x4', 'input2:4x3x2x1']) + assert args.input_shapes == dict(input1=[1, 2, 3, 4], input2=[4, 3, 2, 1]) + + args = parser.parse_args( + ['--min-shapes', 'input1:1x2x3x4, input2:4x3x?x?']) + assert args.min_shapes == dict( + input1=[1, 2, 3, 4], input2=[4, 3, None, None]) + + # input placeholder + with pytest.raises(ValueError): + parser.parse_args(['--input-shapes', 'input1:1x2x3x4, input2:4x3x?x?']) + + # duplicate assign + with pytest.raises(NameError): + parser.parse_args(['--input-shapes', 'input1:1x2x3x4, input1:4x3x2x1']) + + args = parser.parse_args(['--input-shapes', '1x2x3x4']) + assert args.input_shapes == {None: [1, 2, 3, 4]} + + +def test_fix_param(): + + params = BaseBackendParam() + params.input_names = ['input'] + params.input_shapes = {None: [1, 2, 3, 4]} + params.min_shapes = {None: [1, 2, 3, 4]} + params.max_shapes = {None: [1, 2, 3, 4]} + + params.fix_param() + assert params.input_shapes == {'input': [1, 2, 3, 4]} + assert params.min_shapes == {'input': [1, 2, 3, 4]} + assert params.max_shapes == {'input': [1, 2, 3, 4]} + + params.min_shapes = {None: [1, 2, None, None]} + params.max_shapes = {None: [1, 2, None, None]} + params.fix_param() + assert params.min_shapes == {'input': [1, 2, 3, 4]} + assert params.max_shapes == {'input': [1, 2, 3, 4]} + + params = BaseBackendParam(input_shapes={'input': [1, 2, 3, 4]}) + params.fix_param() + assert params.input_names == ['input'] + + # input names + with pytest.raises(ValueError): + params = BaseBackendParam(input_names=['input', 'input']) + params.fix_param() + + with pytest.raises(ValueError): + params = BaseBackendParam(input_names=['input', None]) + params.fix_param() + + # none name error + with pytest.raises(ValueError): + params = BaseBackendParam( + input_names=['input', 'invalid'], + input_shapes={None: [1, 2, 3, 4]}) + params.fix_param() + + # fill none name error + with pytest.raises(ValueError): + params = BaseBackendParam(input_shapes={None: [1, 2, 3, 4]}) + params.fix_param() + + with pytest.raises(ValueError): + params = BaseBackendParam( + input_names=['input'], + input_shapes={ + None: [1, 2, 3, 4], + 'input': [1, 2, 3, 4] + }) + params.fix_param() + + with pytest.raises(ValueError): + params = BaseBackendParam( + input_names=[None], input_shapes={None: [1, 2, 3, 4]}) + params.fix_param() + + # shape type error + with pytest.raises(TypeError): + params = BaseBackendParam(input_names=['input'], input_shapes=0) + params.fix_param() + + with pytest.raises(TypeError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, min_shapes=0) + params.fix_param() + + with pytest.raises(TypeError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, max_shapes=0) + params.fix_param() + + # placeholder error + with pytest.raises(ValueError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, + min_shapes={'input': [1, 2, 3, None, None]}) + params.fix_param() + + with pytest.raises(ValueError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, + max_shapes={'input': [1, 2, 3, None, None]}) + params.fix_param() + + +def test_check_param(): + + params = BaseBackendParam( + input_names=['input'], input_shapes={None: [1, 2, 3, 4]}) + params.check_param() + + # input shapes != min/max shape + with pytest.raises(ValueError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, + min_shapes={'input': [1, 2, 3, 4, 5]}) + params.check_param() + + # different input names + with pytest.raises(NameError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, + min_shapes={'input1': [1, 2, 3, 4]}) + params.check_param() + + with pytest.raises(NameError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, + max_shapes={'input1': [1, 2, 3, 4]}) + params.check_param() + + # shape type error + with pytest.raises(TypeError): + params = BaseBackendParam(input_shapes=0) + params.check_param() + + with pytest.raises(TypeError): + params = BaseBackendParam(min_shapes=0) + params.check_param() + + with pytest.raises(TypeError): + params = BaseBackendParam(max_shapes=0) + params.check_param() + + # shape length error + with pytest.raises(ValueError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, + max_shapes={'input': [1, 2, 3, 4, 5]}) + params.check_param() + + # shape value error + with pytest.raises(ValueError): + params = BaseBackendParam( + input_shapes={'input': [1, 2, 3, 4]}, + max_shapes={'input': [1, 1, 1, 1]}) + params.check_param() diff --git a/tests/test_backend/test_ncnn.py b/tests/test_backend/test_ncnn.py new file mode 100644 index 0000000000..1ab48c1f58 --- /dev/null +++ b/tests/test_backend/test_ncnn.py @@ -0,0 +1,91 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.ncnn import NCNNManager as backend_mgr +from mmdeploy.backend.ncnn import NCNNParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.param' +_bin_extension = '.bin' + + +class TestBackendParam: + + def test_get_model_files(self): + param = NCNNParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + assert param.bin_name == 'tmp' + _bin_extension + + assert param.get_model_files() == ('tmp' + _extension, + 'tmp' + _bin_extension) + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_dynamic_2i2o): + yield onnx_model_dynamic_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model): + with TemporaryDirectory() as tmp_dir: + param_path = osp.join(tmp_dir, 'tmp' + _extension) + bin_path = osp.join(tmp_dir, 'tmp' + _bin_extension) + backend_mgr.to_backend(onnx_model, param_path, bin_path) + + yield param_path, bin_path + + def test_to_backend(self, backend_model): + for path in backend_model: + assert osp.exists(path) + + def test_to_backend_from_param(self, onnx_model): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param(work_dir=work_dir, file_name='tmp') + backend_mgr.to_backend_from_param(onnx_model, param) + + param_path, bin_path = param.get_model_files() + assert osp.exists(param_path) + assert osp.exists(bin_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, + assert_forward): + wrapper = backend_mgr.build_wrapper(*backend_model) + assert_forward(wrapper, inputs, outputs) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + assert_forward): + param_path, bin_path = backend_model + param = backend_mgr.build_param( + work_dir='', file_name=param_path, bin_name=bin_path) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs) + + def test_parse_args(self, onnx_model): + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + _extension + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, param_name)) + assert osp.exists(osp.join(work_dir, 'tmp' + _bin_extension)) diff --git a/tests/test_backend/test_onnxruntime.py b/tests/test_backend/test_onnxruntime.py new file mode 100644 index 0000000000..3a9fb446a1 --- /dev/null +++ b/tests/test_backend/test_onnxruntime.py @@ -0,0 +1,45 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os.path as osp + +import pytest + +from mmdeploy.backend.onnxruntime import ONNXRuntimeManager as backend_mgr + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_dynamic_2i2o): + yield onnx_model_dynamic_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model_dynamic_2i2o): + yield onnx_model_dynamic_2i2o + + def test_to_backend_from_param(self, tmp_path, backend_model): + save_path = str(tmp_path / 'tmp.onnx') + param = backend_mgr.build_param(work_dir='', file_name=save_path) + backend_mgr.to_backend_from_param(backend_model, param) + assert osp.exists(save_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, + assert_forward): + wrapper = backend_mgr.build_wrapper(backend_model, 'cpu') + assert_forward(wrapper, inputs, outputs) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + assert_forward): + param = backend_mgr.build_param(work_dir='', file_name=backend_model) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs) diff --git a/tests/test_backend/test_openvino.py b/tests/test_backend/test_openvino.py new file mode 100644 index 0000000000..fa5a7ff821 --- /dev/null +++ b/tests/test_backend/test_openvino.py @@ -0,0 +1,112 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.openvino import OpenVINOManager as backend_mgr +from mmdeploy.backend.openvino import OpenVINOParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.xml' +_bin_extension = '.bin' + + +class TestBackendParam: + + def test_get_model_files(self): + param = OpenVINOParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + assert param.bin_name == 'tmp' + _bin_extension + + assert param.get_model_files() == ('tmp' + _extension, + 'tmp' + _bin_extension) + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_2i): + yield input_shape_dict_2i + + @pytest.fixture(scope='class') + def output_names(self, output_names_2i2o): + yield output_names_2i2o + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_dynamic_2i2o): + yield onnx_model_dynamic_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model, input_shape_dict, output_names): + with TemporaryDirectory() as tmp_dir: + param_path = osp.join(tmp_dir, 'tmp' + _extension) + bin_path = osp.join(tmp_dir, 'tmp' + _bin_extension) + backend_mgr.to_backend( + onnx_model, + param_path, + input_info=input_shape_dict, + output_names=output_names) + + yield param_path, bin_path + + def test_to_backend(self, backend_model): + for path in backend_model: + assert osp.exists(path) + + def test_to_backend_from_param(self, onnx_model): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param(work_dir=work_dir, file_name='tmp') + backend_mgr.to_backend_from_param(onnx_model, param) + + param_path, bin_path = param.get_model_files() + assert osp.exists(param_path) + assert osp.exists(bin_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, + assert_forward): + wrapper = backend_mgr.build_wrapper(*backend_model) + assert_forward(wrapper, inputs, outputs) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + assert_forward): + param_path, bin_path = backend_model + param = backend_mgr.build_param( + work_dir='', file_name=param_path, bin_name=bin_path) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs) + + def test_parse_args(self, onnx_model, input_shape_dict, output_names): + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + _extension + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + args += ['--output-names', *output_names] + args += ['--input-shapes', input_shapes] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, param_name)) + assert osp.exists(osp.join(work_dir, 'tmp' + _bin_extension)) diff --git a/tests/test_backend/test_pplnn.py b/tests/test_backend/test_pplnn.py new file mode 100644 index 0000000000..8d948ccc9e --- /dev/null +++ b/tests/test_backend/test_pplnn.py @@ -0,0 +1,107 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.pplnn import PPLNNManager as backend_mgr +from mmdeploy.backend.pplnn import PPLNNParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.onnx' +_json_extension = '.json' + + +class TestBackendParam: + + def test_get_model_files(self): + param = PPLNNParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + assert param.algo_name == 'tmp' + _json_extension + + assert param.get_model_files() == ('tmp' + _extension, + 'tmp' + _json_extension) + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_2i): + yield input_shape_dict_2i + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_dynamic_2i2o): + yield onnx_model_dynamic_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model, input_shape_dict): + with TemporaryDirectory() as tmp_dir: + param_path = osp.join(tmp_dir, 'tmp' + _extension) + algo_path = osp.join(tmp_dir, 'tmp' + _json_extension) + backend_mgr.to_backend( + onnx_model, + param_path, + algo_path, + input_shapes=input_shape_dict) + + yield param_path, algo_path + + def test_to_backend(self, backend_model): + assert osp.exists(backend_model[0]) + + def test_to_backend_from_param(self, onnx_model, input_shape_dict): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_shapes=input_shape_dict) + backend_mgr.to_backend_from_param(onnx_model, param) + + param_path, _ = param.get_model_files() + assert osp.exists(param_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, + assert_forward): + wrapper = backend_mgr.build_wrapper(*backend_model) + assert_forward(wrapper, inputs, outputs) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + assert_forward): + param_path, algo_path = backend_model + param = backend_mgr.build_param( + work_dir='', file_name=param_path, algo_name=algo_path) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs) + + def test_parse_args(self, onnx_model, input_shape_dict): + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + with TemporaryDirectory() as work_dir: + file_name = 'tmp' + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', file_name] + args += ['--input-shapes', input_shapes] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, file_name + '.onnx')) diff --git a/tests/test_backend/test_rknn.py b/tests/test_backend/test_rknn.py new file mode 100644 index 0000000000..23285ff102 --- /dev/null +++ b/tests/test_backend/test_rknn.py @@ -0,0 +1,160 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.rknn import RKNNManager as backend_mgr +from mmdeploy.backend.rknn import RKNNParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +device = 'rk1126' +_extension = '.rknn' + + +class TestBackendParam: + + def test_get_model_files(self): + param = RKNNParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + + assert param.get_model_files() == 'tmp' + _extension + + def test_add_argument(self): + parser = argparse.ArgumentParser() + RKNNParam.add_argument( + parser, 'mean_values', dtype=str, default=None, desc='') + + assert parser.parse_args(['--mean-values', + '3.2,-2,1']).mean_values == { + None: (3.2, -2.0, 1.0) + } + assert parser.parse_args(['--mean-values', + 'input:3.2,-2,1']).mean_values == { + 'input': (3.2, -2.0, 1.0) + } + assert parser.parse_args( + ['--mean-values', 'input1:3.2,-2,1', + 'input2:3.2,-2,1']).mean_values == { + 'input1': (3.2, -2.0, 1.0), + 'input2': (3.2, -2.0, 1.0) + } + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def input_names(self, input_names_2i): + yield input_names_2i + + @pytest.fixture(scope='class') + def output_names(self, output_names_2i2o): + yield output_names_2i2o + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_2i): + yield input_shape_dict_2i + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_static_2i2o): + yield onnx_model_static_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model, input_names, output_names, + input_shape_dict): + from mmdeploy.backend.rknn.onnx2rknn import RKNNConfig + with TemporaryDirectory() as tmp_dir: + model_path = osp.join(tmp_dir, 'tmp' + _extension) + backend_mgr.to_backend( + onnx_model, + model_path, + input_names, + output_names, + input_shapes=input_shape_dict, + rknn_config=RKNNConfig( + mean_values=[(0, 0, 0), (0, 0, 0)], + std_values=[(1, 1, 1), (1, 1, 1)], + target_platform=device), + do_quantization=False) + + yield model_path + + def test_to_backend(self, backend_model): + assert osp.exists(backend_model[0]) + + def test_to_backend_from_param(self, onnx_model, input_names, output_names, + input_shape_dict): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_names=input_names, + output_names=output_names, + input_shapes=input_shape_dict, + device=device, + mean_values=dict(x=(0, 0, 0), y=(0, 0, 0)), + std_values=dict(x=(1, 1, 1), y=(1, 1, 1))) + backend_mgr.to_backend_from_param(onnx_model, param) + + param_path = param.get_model_files() + assert osp.exists(param_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, input_names, + output_names, assert_forward): + + wrapper = backend_mgr.build_wrapper( + backend_model, + device, + input_names, + output_names, + ) + assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + input_names, output_names, + assert_forward): + param = backend_mgr.build_param( + work_dir='', + device=device, + file_name=backend_model, + input_names=input_names, + output_names=output_names) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + def test_parse_args(self, onnx_model, input_names, output_names, + input_shape_dict): + + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + _extension + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + args += ['--input-names', *input_names] + args += ['--output-names', *output_names] + args += ['--input-shapes', input_shapes] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, param_name)) diff --git a/tests/test_backend/test_snpe.py b/tests/test_backend/test_snpe.py new file mode 100644 index 0000000000..d2a6b330c1 --- /dev/null +++ b/tests/test_backend/test_snpe.py @@ -0,0 +1,72 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.snpe import SNPEManager as backend_mgr +from mmdeploy.backend.snpe import SNPEParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.dlc' + + +class TestBackendParam: + + def test_get_model_files(self): + param = SNPEParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + + assert param.get_model_files() == 'tmp' + _extension + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def ir_model(self, onnx_model_static_2i2o): + yield onnx_model_static_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, ir_model): + with TemporaryDirectory() as tmp_dir: + model_path = osp.join(tmp_dir, 'tmp' + _extension) + backend_mgr.to_backend(ir_model, model_path) + + yield model_path + + def test_to_backend(self, backend_model): + assert osp.exists(backend_model) + + def test_to_backend_from_param(self, ir_model): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param(work_dir=work_dir, file_name='tmp') + backend_mgr.to_backend_from_param(ir_model, param) + + param_path = param.get_model_files() + assert osp.exists(param_path) + + def test_parse_args(self, ir_model): + + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + _extension + # make args + args = ['convert'] + args += ['--onnx-path', ir_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, param_name)) diff --git a/tests/test_backend/test_tensorrt.py b/tests/test_backend/test_tensorrt.py new file mode 100644 index 0000000000..dc2aa75fa2 --- /dev/null +++ b/tests/test_backend/test_tensorrt.py @@ -0,0 +1,139 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.tensorrt import TensorRTManager as backend_mgr +from mmdeploy.backend.tensorrt import TensorRTParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.engine' + + +class TestBackendParam: + + def test_get_model_files(self): + param = TensorRTParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + _extension + + assert param.get_model_files() == 'tmp' + _extension + + def test_check_param(self): + with pytest.raises(ValueError): + param = TensorRTParam(int8_mode=True, int8_algorithm='invalid') + param.check_param() + + def test_parse_args(self): + parser = argparse.ArgumentParser() + TensorRTParam.add_arguments(parser) + args = parser.parse_args([ + '--fp16-mode', '--int8-mode', '--int8-algorithm', 'maxmin', + '--max-workspace-size', '1024' + ]) + + assert getattr(args, 'device', None) == 'cuda' + assert getattr(args, 'fp16_mode', None) is True + assert getattr(args, 'int8_mode', None) is True + assert getattr(args, 'int8_algorithm', None) == 'maxmin' + assert getattr(args, 'max_workspace_size', None) == 1024 + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + torch = pytest.importorskip('torch') + if not torch.cuda.is_available(): + pytest.skip('torch cuda is not available') + cuda_inputs = dict((k, v.cuda()) for k, v in input_dict_2i.items()) + yield cuda_inputs + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + torch = pytest.importorskip('torch') + if not torch.cuda.is_available(): + pytest.skip('torch cuda is not available') + cuda_outputs = dict((k, v.cuda()) for k, v in output_dict_2i2o.items()) + yield cuda_outputs + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_2i): + yield input_shape_dict_2i + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_dynamic_2i2o): + yield onnx_model_dynamic_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, input_shape_dict, onnx_model_dynamic_2i2o): + from tempfile import NamedTemporaryFile + save_path = NamedTemporaryFile(suffix='.engine').name + backend_mgr.to_backend(onnx_model_dynamic_2i2o, save_path, + input_shape_dict) + yield save_path + + def test_to_backend(self, backend_model): + assert osp.exists(backend_model) + + def test_to_backend_from_param(self, input_shape_dict, onnx_model): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_shapes=input_shape_dict) + backend_mgr.to_backend_from_param(onnx_model, param) + assert osp.exists(param.get_model_files()) + + def test_to_backend_from_param_quanti(self, input_shape_dict, onnx_model, + inputs): + + def _quanti_data(): + yield inputs + + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_shapes=input_shape_dict, + int8_mode=True, + quanti_data=_quanti_data()) + backend_mgr.to_backend_from_param(onnx_model, param) + assert osp.exists(param.get_model_files()) + + def test_build_wrapper(self, backend_model, inputs, outputs, + assert_forward): + wrapper = backend_mgr.build_wrapper(backend_model) + assert_forward(wrapper, inputs, outputs) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + assert_forward): + param = backend_mgr.build_param(work_dir='', file_name=backend_model) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs) + + def test_parse_args(self, onnx_model, input_shape_dict): + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + with TemporaryDirectory() as work_dir: + + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', 'tmp'] + args += ['--input-shapes', input_shapes] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + + assert osp.exists(osp.join(work_dir, 'tmp' + _extension)) diff --git a/tests/test_backend/test_torchjit.py b/tests/test_backend/test_torchjit.py new file mode 100644 index 0000000000..a66fc8ba17 --- /dev/null +++ b/tests/test_backend/test_torchjit.py @@ -0,0 +1,59 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os.path as osp + +import pytest + +from mmdeploy.backend.torchscript import TorchScriptManager as backend_mgr + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def input_names(self, input_names_2i): + yield input_names_2i + + @pytest.fixture(scope='class') + def output_names(self, output_names_2i2o): + yield output_names_2i2o + + @pytest.fixture(scope='class') + def torchscript_model(self, torchscript_model2i2o): + yield torchscript_model2i2o + + @pytest.fixture(scope='class') + def backend_model(self, torchscript_model): + yield torchscript_model + + def test_to_backend_from_param(self, tmp_path, backend_model): + save_path = str(tmp_path / 'tmp.pth') + param = backend_mgr.build_param(work_dir='', file_name=save_path) + backend_mgr.to_backend_from_param(backend_model, param) + assert osp.exists(save_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, input_names, + output_names, assert_forward): + wrapper = backend_mgr.build_wrapper(backend_model, input_names, + output_names) + assert_forward(wrapper, inputs, outputs) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + input_names, output_names, + assert_forward): + param = backend_mgr.build_param( + work_dir='', + file_name=backend_model, + input_names=input_names, + output_names=output_names) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs) diff --git a/tests/test_backend/test_tvm.py b/tests/test_backend/test_tvm.py new file mode 100644 index 0000000000..9793a8d700 --- /dev/null +++ b/tests/test_backend/test_tvm.py @@ -0,0 +1,137 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import pytest + +from mmdeploy.backend.tvm import TVMManager as backend_mgr +from mmdeploy.backend.tvm import TVMParam, get_library_ext + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + + +class TestBackendParam: + + def test_get_model_files(self): + param = TVMParam(work_dir='', file_name='tmp') + assert param.file_name == 'tmp' + get_library_ext() + assert param.vm_name == 'tmp.vm' + + assert param.get_model_files() == ('tmp' + get_library_ext(), 'tmp.vm') + + def test_add_argument(self): + parser = argparse.ArgumentParser() + TVMParam.add_argument( + parser, 'dtypes', dtype=str, default=None, desc='') + + assert parser.parse_args(['--dtypes', 'float32']).dtypes == { + None: 'float32' + } + assert parser.parse_args(['--dtypes', 'input:float32']).dtypes == { + 'input': 'float32' + } + assert parser.parse_args( + ['--dtypes', 'input1:float32', 'input2:float32']).dtypes == { + 'input1': 'float32', + 'input2': 'float32' + } + assert parser.parse_args(['--dtypes', + 'input1:float32,input2:float32']).dtypes == { + 'input1': 'float32', + 'input2': 'float32' + } + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_2i): + yield input_dict_2i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_2i2o): + yield output_dict_2i2o + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_2i): + yield input_shape_dict_2i + + @pytest.fixture(scope='class') + def input_dtypes(self): + yield dict(x='float32', y='float32') + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_dynamic_2i2o): + yield onnx_model_dynamic_2i2o + + @pytest.fixture(scope='class') + def output_names(self, output_names_2i2o): + yield output_names_2i2o + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model, input_shape_dict, input_dtypes): + with TemporaryDirectory() as tmp_dir: + output_path = osp.join(tmp_dir, 'tmp' + get_library_ext()) + backend_mgr.to_backend( + onnx_model, + output_path, + input_shapes=input_shape_dict, + dtypes=input_dtypes) + + yield output_path + + def test_to_backend(self, backend_model): + assert osp.exists(backend_model) + + def test_to_backend_from_param(self, onnx_model, input_shape_dict, + input_dtypes): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_shapes=input_shape_dict, + dtypes=input_dtypes) + backend_mgr.to_backend_from_param(onnx_model, param) + + param_path, _ = param.get_model_files() + assert osp.exists(param_path) + + def test_build_wrapper(self, backend_model, inputs, outputs, output_names, + assert_forward): + wrapper = backend_mgr.build_wrapper( + backend_model, output_names=output_names) + assert_forward(wrapper, inputs, outputs) + + def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + output_names, assert_forward): + param = backend_mgr.build_param( + work_dir='', file_name=backend_model, output_names=output_names) + wrapper = backend_mgr.build_wrapper_from_param(param) + assert_forward(wrapper, inputs, outputs) + + def test_parse_args(self, onnx_model, input_shape_dict, input_dtypes): + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + dtypes = ','.join(f'{k}:{v}' for k, v in input_dtypes.items()) + + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + get_library_ext() + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + args += ['--input-shapes', input_shapes] + args += ['--dtypes', dtypes] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists(osp.join(work_dir, param_name)) diff --git a/tests/test_backend/test_vacc.py b/tests/test_backend/test_vacc.py new file mode 100644 index 0000000000..6bf1a212e6 --- /dev/null +++ b/tests/test_backend/test_vacc.py @@ -0,0 +1,179 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import os.path as osp +from tempfile import TemporaryDirectory + +import numpy as np +import pytest + +from mmdeploy.backend.vacc import VACCManager as backend_mgr +from mmdeploy.backend.vacc import VACCParam + +if not backend_mgr.is_available(): + pytest.skip('backend not available', allow_module_level=True) + +_extension = '.so' +_json_extension = '.json' +_param_extension = '.params' + + +class TestBackendParam: + + def test_get_model_files(self): + param = VACCParam(work_dir='', file_name='tmp') + + assert param.get_model_files() == [ + 'tmp-fp16/tmp' + _extension, 'tmp-fp16/tmp' + _json_extension, + 'tmp-fp16/tmp' + _param_extension + ] + + +def half_to_uint16(x): + """Convert a np.float16 number to a uint16 represented.""" + return int(np.frombuffer(np.array(x, dtype=np.float16), dtype=np.uint16)) + + +class TestManager: + + @pytest.fixture(scope='class') + def inputs(self, input_dict_1i): + yield input_dict_1i + + @pytest.fixture(scope='class') + def outputs(self, output_dict_1i1o): + yield output_dict_1i1o + + @pytest.fixture(scope='class') + def input_names(self, input_names_1i): + yield input_names_1i + + @pytest.fixture(scope='class') + def output_names(self, output_names_1i1o): + yield output_names_1i1o + + @pytest.fixture(scope='class') + def input_shape_dict(self, input_shape_dict_1i): + yield input_shape_dict_1i + + @pytest.fixture(scope='class') + def onnx_model(self, onnx_model_static_1i1o): + yield onnx_model_static_1i1o + + @pytest.fixture(scope='class') + def work_dir(self): + with TemporaryDirectory() as tmp_dir: + yield tmp_dir + + @pytest.fixture(scope='class') + def backend_model(self, onnx_model, input_shape_dict, work_dir): + model_name = 'tmp' + outputs = VACCParam( + work_dir=work_dir, file_name=model_name, quant_mode='fp16') + backend_mgr.to_backend( + onnx_model, + work_dir, + model_name, + input_shapes=input_shape_dict, + quant_mode='fp16') + + yield outputs.get_model_files() + + @pytest.fixture(scope='class') + def dummy_vdsp_params_info(self): + mean = half_to_uint16(0) + std = half_to_uint16(255) + return dict( + vdsp_op_type=300, + iimage_format=5000, + iimage_width=8, + iimage_height=8, + iimage_width_pitch=8, + iimage_height_pitch=8, + short_edge_threshold=8, + resize_type=1, + color_cvt_code=2, + color_space=0, + crop_size=8, + meanr=mean, + meang=mean, + meanb=mean, + stdr=std, + stdg=std, + stdb=std, + norma_type=3) + + def test_to_backend(self, backend_model): + for file in backend_model: + assert osp.exists(file) + + def test_to_backend_from_param(self, onnx_model, input_shape_dict): + with TemporaryDirectory() as work_dir: + param = backend_mgr.build_param( + work_dir=work_dir, + file_name='tmp', + input_shapes=input_shape_dict) + backend_mgr.to_backend_from_param(onnx_model, param) + + model_files = param.get_model_files() + for file in model_files: + assert osp.exists(file) + + # TODO: Enable the test after vdsp parameter available + # def test_build_wrapper(self, backend_model, inputs, outputs, + # dummy_vdsp_params_info, output_names, + # assert_forward): + + # wrapper = backend_mgr.build_wrapper( + # *backend_model, + # dummy_vdsp_params_info, + # output_names, + # ) + # assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + # def test_build_wrapper_from_param(self, backend_model, inputs, outputs, + # dummy_vdsp_params_info, output_names, + # work_dir, assert_forward): + # for file in backend_model: + # assert osp.exists(file) + + # param = backend_mgr.build_param( + # work_dir=work_dir, + # file_name='tmp', + # output_names=output_names, + # vdsp_params_info=dummy_vdsp_params_info) + # wrapper = backend_mgr.build_wrapper_from_param(param) + # assert_forward(wrapper, inputs, outputs, rtol=1e-3, atol=1e-3) + + def test_parse_args(self, onnx_model, output_names, input_shape_dict): + + # make input shapes + input_shapes = [] + for name, shape in input_shape_dict.items(): + shape = 'x'.join(str(i) for i in shape) + input_shapes.append(f'{name}:{shape}') + input_shapes = ','.join(input_shapes) + + with TemporaryDirectory() as work_dir: + param_name = 'tmp' + quant_mode = 'fp16' + # make args + args = ['convert'] + args += ['--onnx-path', onnx_model] + args += ['--work-dir', work_dir] + args += ['--file-name', param_name] + args += ['--output-names', *output_names] + args += ['--input-shapes', input_shapes] + args += ['--quant-mode', quant_mode] + + parser = argparse.ArgumentParser() + with backend_mgr.parse_args(parser, args=args): + pass + assert osp.exists( + osp.join(work_dir, param_name + '-' + quant_mode, + param_name + _extension)) + assert osp.exists( + osp.join(work_dir, param_name + '-' + quant_mode, + param_name + _json_extension)) + assert osp.exists( + osp.join(work_dir, param_name + '-' + quant_mode, + param_name + _param_extension)) diff --git a/tests/test_backend/test_wrapper.py b/tests/test_backend/test_wrapper.py deleted file mode 100644 index c2d21f5c14..0000000000 --- a/tests/test_backend/test_wrapper.py +++ /dev/null @@ -1,219 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -import os.path as osp -import subprocess -import tempfile - -import mmengine -import pytest -import torch -import torch.nn as nn - -from mmdeploy.utils.constants import Backend -from mmdeploy.utils.test import check_backend - -onnx_file = tempfile.NamedTemporaryFile(suffix='.onnx').name -ts_file = tempfile.NamedTemporaryFile(suffix='.pt').name -test_img = torch.rand(1, 3, 8, 8) -output_names = ['output'] -input_names = ['input'] -target_platform = 'rk3588' # rknn pre-compiled model need device - - -@pytest.mark.skip(reason='This a not test class but a utility class.') -class TestModel(nn.Module): - - def __init__(self): - super().__init__() - - def forward(self, x): - return x + test_img - - -model = TestModel().eval() - - -@pytest.fixture(autouse=True, scope='module') -def generate_onnx_file(): - with torch.no_grad(): - torch.onnx.export( - model, - test_img, - onnx_file, - output_names=output_names, - input_names=input_names, - keep_initializers_as_inputs=True, - do_constant_folding=True, - verbose=False, - opset_version=11, - dynamic_axes=None) - - -@pytest.fixture(autouse=True, scope='module') -def generate_torchscript_file(): - from mmengine import Config - - backend = Backend.TORCHSCRIPT.value - deploy_cfg = Config({'backend_config': dict(type=backend)}) - - from mmdeploy.apis.torch_jit import trace - context_info = dict(deploy_cfg=deploy_cfg) - output_prefix = osp.splitext(ts_file)[0] - - example_inputs = torch.rand(1, 3, 8, 8) - trace( - model, - example_inputs, - output_path_prefix=output_prefix, - backend=backend, - context_info=context_info) - - -def ir2backend(backend, onnx_file, ts_file): - if backend == Backend.TENSORRT: - from mmdeploy.backend.tensorrt import from_onnx - backend_file = tempfile.NamedTemporaryFile(suffix='.engine').name - from_onnx( - onnx_file, - osp.splitext(backend_file)[0], { - 'input': { - 'min_shape': [1, 3, 8, 8], - 'opt_shape': [1, 3, 8, 8], - 'max_shape': [1, 3, 8, 8] - } - }) - return backend_file - elif backend == Backend.ONNXRUNTIME: - return onnx_file - elif backend == Backend.PPLNN: - from mmdeploy.apis.pplnn import from_onnx - output_file_prefix = tempfile.NamedTemporaryFile().name - from_onnx(onnx_file, output_file_prefix=output_file_prefix) - algo_file = output_file_prefix + '.json' - output_file = output_file_prefix + '.onnx' - return output_file, algo_file - elif backend == Backend.NCNN: - from mmdeploy.backend.ncnn.init_plugins import get_onnx2ncnn_path - onnx2ncnn_path = get_onnx2ncnn_path() - param_file = tempfile.NamedTemporaryFile(suffix='.param').name - bin_file = tempfile.NamedTemporaryFile(suffix='.bin').name - subprocess.call([onnx2ncnn_path, onnx_file, param_file, bin_file]) - return param_file, bin_file - elif backend == Backend.OPENVINO: - from mmdeploy.apis.openvino import from_onnx, get_output_model_file - backend_dir = tempfile.TemporaryDirectory().name - backend_file = get_output_model_file(onnx_file, backend_dir) - input_info = {'input': test_img.shape} - output_names = ['output'] - work_dir = backend_dir - from_onnx(onnx_file, work_dir, input_info, output_names) - return backend_file - elif backend == Backend.RKNN: - from mmdeploy.apis.rknn import onnx2rknn - rknn_file = onnx_file.replace('.onnx', '.rknn') - deploy_cfg = mmengine.Config( - dict( - backend_config=dict( - type='rknn', - common_config=dict(target_platform=target_platform), - quantization_config=dict( - do_quantization=False, dataset=None), - input_size_list=[[3, 8, 8]]))) - onnx2rknn(onnx_file, rknn_file, deploy_cfg) - return rknn_file - elif backend == Backend.ASCEND: - from mmdeploy.apis.ascend import from_onnx - backend_dir = tempfile.TemporaryDirectory().name - work_dir = backend_dir - file_name = osp.splitext(osp.split(onnx_file)[1])[0] - backend_file = osp.join(work_dir, file_name + '.om') - model_inputs = mmengine.Config( - dict(input_shapes=dict(input=test_img.shape))) - from_onnx(onnx_file, work_dir, model_inputs) - return backend_file - elif backend == Backend.TVM: - from mmdeploy.backend.tvm import from_onnx, get_library_ext - ext = get_library_ext() - lib_file = tempfile.NamedTemporaryFile(suffix=ext).name - shape = {'input': test_img.shape} - dtype = {'input': 'float32'} - target = 'llvm' - tuner_dict = dict(type='DefaultTuner', target=target) - from_onnx( - onnx_file, lib_file, shape=shape, dtype=dtype, tuner=tuner_dict) - assert osp.exists(lib_file) - return lib_file - elif backend == Backend.TORCHSCRIPT: - return ts_file - elif backend == Backend.COREML: - output_names = ['output'] - from mmdeploy.backend.coreml.torchscript2coreml import ( - from_torchscript, get_model_suffix) - backend_dir = tempfile.TemporaryDirectory().name - work_dir = backend_dir - torchscript_name = osp.splitext(osp.split(ts_file)[1])[0] - output_file_prefix = osp.join(work_dir, torchscript_name) - convert_to = 'mlprogram' - from_torchscript( - ts_file, - output_file_prefix, - input_names=input_names, - output_names=output_names, - input_shapes=dict( - input=dict( - min_shape=[1, 3, 8, 8], - default_shape=[1, 3, 8, 8], - max_shape=[1, 3, 8, 8])), - convert_to=convert_to) - suffix = get_model_suffix(convert_to) - return output_file_prefix + suffix - else: - raise NotImplementedError( - f'Convert for {backend.value} has not been implemented.') - - -def create_wrapper(backend, model_files): - from mmdeploy.backend.base import get_backend_manager - backend_mgr = get_backend_manager(backend.value) - deploy_cfg = None - if isinstance(model_files, str): - model_files = [model_files] - elif backend == Backend.RKNN: - deploy_cfg = dict( - backend_config=dict( - common_config=dict(target_platform=target_platform))) - return backend_mgr.build_wrapper( - model_files, - input_names=input_names, - output_names=output_names, - deploy_cfg=deploy_cfg) - - -def run_wrapper(backend, wrapper, input): - if backend == Backend.TENSORRT: - input = input.cuda() - - results = wrapper({'input': input}) - - if backend != Backend.RKNN: - results = results['output'] - - results = results.detach().cpu() - - return results - - -ALL_BACKEND = list(Backend) -ALL_BACKEND.remove(Backend.DEFAULT) -ALL_BACKEND.remove(Backend.PYTORCH) -ALL_BACKEND.remove(Backend.SDK) - - -@pytest.mark.parametrize('backend', ALL_BACKEND) -def test_wrapper(backend): - check_backend(backend, True) - model_files = ir2backend(backend, onnx_file, ts_file) - assert model_files is not None - wrapper = create_wrapper(backend, model_files) - assert wrapper is not None - results = run_wrapper(backend, wrapper, test_img) - assert results is not None diff --git a/tests/test_backend/utils.py b/tests/test_backend/utils.py new file mode 100644 index 0000000000..ef101fec61 --- /dev/null +++ b/tests/test_backend/utils.py @@ -0,0 +1 @@ +# Copyright (c) OpenMMLab. All rights reserved. diff --git a/tests/test_codebase/test_mmaction/test_video_recognition.py b/tests/test_codebase/test_mmaction/test_video_recognition.py index ef61df0db4..b31a28401f 100644 --- a/tests/test_codebase/test_mmaction/test_video_recognition.py +++ b/tests/test_codebase/test_mmaction/test_video_recognition.py @@ -5,7 +5,6 @@ import torch from mmengine import Config -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -39,18 +38,15 @@ video = 'tests/test_codebase/test_mmaction/data/video/demo.mp4' -@pytest.fixture +@pytest.fixture(scope='class') def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'output': torch.rand(1, num_classes), - }) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(outputs={ + 'output': torch.rand(1, num_classes), + }) - yield task_processor.build_backend_model(['']) - - wrapper.recover() + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmaction/test_video_recognition_model.py b/tests/test_codebase/test_mmaction/test_video_recognition_model.py index 5908ae7639..4e7b222ea6 100644 --- a/tests/test_codebase/test_mmaction/test_video_recognition_model.py +++ b/tests/test_codebase/test_mmaction/test_video_recognition_model.py @@ -4,7 +4,6 @@ import torch from mmengine import Config -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase, load_config from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -21,41 +20,35 @@ @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - - # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'outputs': torch.rand(1, 400), - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config({'onnx_config': {'output_names': ['outputs']}}) - model_cfg_path = 'tests/test_codebase/test_mmaction/data/model.py' - model_cfg = load_config(model_cfg_path)[0] - + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper from mmdeploy.codebase.mmaction.deploy.video_recognition_model import \ End2EndModel - cls.end2end_model = End2EndModel( - Backend.ONNXRUNTIME, [''], - device='cpu', - deploy_cfg=deploy_cfg, - model_cfg=model_cfg) - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + # simplify backend inference + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'outputs': torch.rand(1, 400), + } + wrapper.set(outputs=outputs) + deploy_cfg = Config({'onnx_config': {'output_names': ['outputs']}}) + model_cfg_path = 'tests/test_codebase/test_mmaction/data/model.py' + model_cfg = load_config(model_cfg_path)[0] + + yield End2EndModel( + Backend.ONNXRUNTIME, [''], + device='cpu', + deploy_cfg=deploy_cfg, + model_cfg=model_cfg) - def test_forward(self): + def test_forward(self, end2end_model): inputs = torch.rand(1, 3, 3, IMAGE_SIZE, IMAGE_SIZE) from mmaction.structures import ActionDataSample data_sample = ActionDataSample( metainfo=dict(img_shape=(IMAGE_SIZE, IMAGE_SIZE))) - results = self.end2end_model.forward( - inputs, [data_sample], mode='predict') + results = end2end_model.forward(inputs, [data_sample], mode='predict') assert results is not None, 'failed to get output using '\ 'End2EndModel' @@ -70,8 +63,7 @@ def test_build_video_recognition_model(): onnx_config=dict(output_names=['outputs']), codebase_config=dict(type='mmaction'))) - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_codebase/test_mmcls/test_classification.py b/tests/test_codebase/test_mmcls/test_classification.py index 3b5c5ca301..899b179099 100644 --- a/tests/test_codebase/test_mmcls/test_classification.py +++ b/tests/test_codebase/test_mmcls/test_classification.py @@ -9,7 +9,6 @@ import torch from mmengine import Config -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -79,18 +78,15 @@ def test_build_pytorch_model(from_mmrazor: Any): assert isinstance(model, BaseClassifier) -@pytest.fixture +@pytest.fixture(scope='class') def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'output': torch.rand(1, num_classes), - }) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(outputs={ + 'output': torch.rand(1, num_classes), + }) - yield task_processor.build_backend_model(['']) - - wrapper.recover() + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmcls/test_classification_model.py b/tests/test_codebase/test_mmcls/test_classification_model.py index 2368cdc165..ee1f773007 100644 --- a/tests/test_codebase/test_mmcls/test_classification_model.py +++ b/tests/test_codebase/test_mmcls/test_classification_model.py @@ -5,7 +5,6 @@ import torch from mmengine import Config -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -23,30 +22,25 @@ @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'outputs': torch.rand(1, NUM_CLASS), - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config({'onnx_config': {'output_names': ['outputs']}}) - - from mmdeploy.codebase.mmcls.deploy.classification_model import \ - End2EndModel - cls.end2end_model = End2EndModel( - Backend.ONNXRUNTIME, [''], device='cpu', deploy_cfg=deploy_cfg) - - @classmethod - def teardown_class(cls): - cls.wrapper.recover() - - def test_forward(self): + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'outputs': torch.rand(1, NUM_CLASS), + } + wrapper.set(outputs=outputs) + deploy_cfg = Config({'onnx_config': {'output_names': ['outputs']}}) + + from mmdeploy.codebase.mmcls.deploy.classification_model import \ + End2EndModel + yield End2EndModel( + Backend.ONNXRUNTIME, [''], device='cpu', deploy_cfg=deploy_cfg) + + def test_forward(self, end2end_model): imgs = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE) from mmcls.structures import ClsDataSample data_sample = ClsDataSample( @@ -54,8 +48,7 @@ def test_forward(self): scale_factor=(1, 1), ori_shape=(IMAGE_SIZE, IMAGE_SIZE), img_shape=(IMAGE_SIZE, IMAGE_SIZE))) - results = self.end2end_model.forward( - imgs, [data_sample], mode='predict') + results = end2end_model.forward(imgs, [data_sample], mode='predict') assert results is not None, 'failed to get output using '\ 'End2EndModel' @@ -63,38 +56,36 @@ def test_forward(self): @backend_checker(Backend.RKNN) class TestRKNNEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - import mmdeploy.backend.rknn as rknn_apis - from mmdeploy.backend.rknn import RKNNWrapper - rknn_apis.__dict__.update({'RKNNWrapper': RKNNWrapper}) + from mmdeploy.backend.rknn.wrapper import RKNNWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(RKNNWrapper) - cls.outputs = [torch.rand(1, 1, IMAGE_SIZE, IMAGE_SIZE)] - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config({ - 'onnx_config': { - 'output_names': ['outputs'] - }, - 'backend_config': { - 'common_config': {} - } - }) - - from mmdeploy.codebase.mmcls.deploy.classification_model import \ - RKNNEnd2EndModel - class_names = ['' for i in range(NUM_CLASS)] - cls.end2end_model = RKNNEnd2EndModel( - Backend.RKNN, [''], - device='cpu', - class_names=class_names, - deploy_cfg=deploy_cfg) - - def test_forward_test(self): + with SwitchBackendWrapper(RKNNWrapper) as wrapper: + outputs = [torch.rand(1, 1, IMAGE_SIZE, IMAGE_SIZE)] + wrapper.set(outputs=outputs) + deploy_cfg = Config({ + 'onnx_config': { + 'output_names': ['outputs'] + }, + 'backend_config': { + 'common_config': {} + } + }) + + from mmdeploy.codebase.mmcls.deploy.classification_model import \ + RKNNEnd2EndModel + class_names = ['' for i in range(NUM_CLASS)] + yield RKNNEnd2EndModel( + Backend.RKNN, [''], + device='cpu', + class_names=class_names, + deploy_cfg=deploy_cfg) + + def test_forward_test(self, end2end_model): imgs = torch.rand(2, 3, IMAGE_SIZE, IMAGE_SIZE) - results = self.end2end_model.forward_test(imgs) + results = end2end_model.forward_test(imgs) assert isinstance(results[0], np.ndarray) @@ -107,8 +98,7 @@ def test_build_classification_model(): onnx_config=dict(output_names=['outputs']), codebase_config=dict(type='mmcls'))) - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_codebase/test_mmdet/test_mmdet_models.py b/tests/test_codebase/test_mmdet/test_mmdet_models.py index c68918e6aa..0d334fccf6 100644 --- a/tests/test_codebase/test_mmdet/test_mmdet_models.py +++ b/tests/test_codebase/test_mmdet/test_mmdet_models.py @@ -206,8 +206,8 @@ def test_multiclass_nms_with_keep_top_k(pre_top_k): test_scores = torch.ones(batch_size, num_boxes, num_classes) model_inputs = {'boxes': test_boxes, 'scores': test_scores} - import mmdeploy.backend.onnxruntime as ort_apis - backend_model = ort_apis.ORTWrapper(onnx_model_path, 'cpu', None) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + backend_model = ORTWrapper(onnx_model_path, 'cpu', None) output = backend_model.forward(model_inputs) output = backend_model.output_to_list(output) dets = output[0] diff --git a/tests/test_codebase/test_mmdet/test_object_detection.py b/tests/test_codebase/test_mmdet/test_object_detection.py index 04525e539e..9d59d55ba1 100644 --- a/tests/test_codebase/test_mmdet/test_object_detection.py +++ b/tests/test_codebase/test_mmdet/test_object_detection.py @@ -11,7 +11,6 @@ from torch.utils.data import DataLoader from torch.utils.data.dataset import Dataset -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -124,18 +123,15 @@ def test_build_pytorch_model(from_mmrazor: Any): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set( - outputs={ - 'dets': torch.rand(1, 10, 5).sort(2).values, - 'labels': torch.randint(0, 10, (1, 10)) - }) - - yield task_processor.build_backend_model(['']) - - wrapper.recover() + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set( + outputs={ + 'dets': torch.rand(1, 10, 5).sort(2).values, + 'labels': torch.randint(0, 10, (1, 10)) + }) + + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmdet/test_object_detection_model.py b/tests/test_codebase/test_mmdet/test_object_detection_model.py index 0091e610e6..4cb6feb61d 100644 --- a/tests/test_codebase/test_mmdet/test_object_detection_model.py +++ b/tests/test_codebase/test_mmdet/test_object_detection_model.py @@ -6,8 +6,6 @@ from mmengine import Config from mmengine.structures import BaseDataElement, InstanceData -import mmdeploy.backend.ncnn as ncnn_apis -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -42,35 +40,28 @@ def assert_forward_results(results, module_name: str = 'model'): @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - # make sure ONNXRuntimeDetector can use ORTWrapper inside itself - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'dets': torch.rand(1, 10, 5), - 'labels': torch.rand(1, 10) - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config( - {'onnx_config': { - 'output_names': ['dets', 'labels'] - }}) - - from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ - End2EndModel - cls.end2end_model = End2EndModel(Backend.ONNXRUNTIME, [''], 'cpu', - deploy_cfg) + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'dets': torch.rand(1, 10, 5), + 'labels': torch.rand(1, 10) + } + wrapper.set(outputs=outputs) + deploy_cfg = Config( + {'onnx_config': { + 'output_names': ['dets', 'labels'] + }}) - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ + End2EndModel + yield End2EndModel(Backend.ONNXRUNTIME, [''], 'cpu', deploy_cfg) - def test_forward(self): + def test_forward(self, end2end_model): imgs = torch.rand(1, 3, 64, 64) img_metas = [ BaseDataElement(metainfo={ @@ -78,12 +69,12 @@ def test_forward(self): 'scale_factor': [1, 1] }) ] - results = self.end2end_model.forward(imgs, img_metas) + results = end2end_model.forward(imgs, img_metas) assert_forward_results(results, 'End2EndModel') - def test_predict(self): + def test_predict(self, end2end_model): imgs = torch.rand(1, 3, 64, 64) - dets, labels = self.end2end_model.predict(imgs) + dets, labels = end2end_model.predict(imgs) assert dets.shape[-1] == 5 assert labels.shape[0] == dets.shape[0] @@ -91,44 +82,37 @@ def test_predict(self): @backend_checker(Backend.ONNXRUNTIME) class TestMaskEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - # make sure ONNXRuntimeDetector can use ORTWrapper inside itself - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference num_classes = 80 num_dets = 10 - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'dets': torch.rand(1, num_dets, 5), - 'labels': torch.randint(num_classes, (1, num_dets)), - 'masks': torch.rand(1, num_dets, 28, 28) - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config({ - 'onnx_config': { - 'output_names': ['dets', 'labels', 'masks'] - }, - 'codebase_config': { - 'post_processing': { - 'export_postprocess_mask': False - } + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'dets': torch.rand(1, num_dets, 5), + 'labels': torch.randint(num_classes, (1, num_dets)), + 'masks': torch.rand(1, num_dets, 28, 28) } - }) - - from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ - End2EndModel - cls.end2end_model = End2EndModel(Backend.ONNXRUNTIME, [''], 'cpu', - deploy_cfg) + wrapper.set(outputs=outputs) + deploy_cfg = Config({ + 'onnx_config': { + 'output_names': ['dets', 'labels', 'masks'] + }, + 'codebase_config': { + 'post_processing': { + 'export_postprocess_mask': False + } + } + }) - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ + End2EndModel + yield End2EndModel(Backend.ONNXRUNTIME, [''], 'cpu', deploy_cfg) - def test_forward(self): + def test_forward(self, end2end_model): imgs = torch.rand(1, 3, 64, 64) img_metas = [ BaseDataElement( @@ -138,7 +122,7 @@ def test_forward(self): 'scale_factor': [1, 1] }) ] - results = self.end2end_model.forward(imgs, img_metas) + results = end2end_model.forward(imgs, img_metas) assert_forward_results(results, 'mask End2EndModel') @@ -239,8 +223,7 @@ def test_build_object_detection_model(partition_type): type=partition_type, partition_cfg=[dict(output_names=[])]) - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: @@ -255,118 +238,109 @@ def test_build_object_detection_model(partition_type): @backend_checker(Backend.NCNN) class TestNCNNEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.ncnn import NCNNWrapper - ncnn_apis.__dict__.update({'NCNNWrapper': NCNNWrapper}) + from mmdeploy.backend.ncnn.wrapper import NCNNWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(NCNNWrapper) - cls.outputs = { - 'output': torch.rand(1, 10, 6), - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config({'onnx_config': {'output_names': ['output']}}) - model_cfg = Config({}) - - from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ - NCNNEnd2EndModel - cls.ncnn_end2end_model = NCNNEnd2EndModel(Backend.NCNN, ['', ''], - 'cpu', model_cfg, deploy_cfg) + with SwitchBackendWrapper(NCNNWrapper) as wrapper: + outputs = { + 'output': torch.rand(1, 10, 6), + } + wrapper.set(outputs=outputs) + deploy_cfg = Config({'onnx_config': {'output_names': ['output']}}) + model_cfg = Config({}) - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ + NCNNEnd2EndModel + yield NCNNEnd2EndModel(Backend.NCNN, ['', ''], 'cpu', model_cfg, + deploy_cfg) @pytest.mark.parametrize('num_det', [10, 0]) - def test_predict(self, num_det): + def test_predict(self, end2end_model, num_det): self.outputs = { 'output': torch.rand(1, num_det, 6), } imgs = torch.rand(1, 3, 64, 64) - results = self.ncnn_end2end_model.predict(imgs) + results = end2end_model.predict(imgs) assert_det_results(results, 'NCNNEnd2EndModel') @backend_checker(Backend.RKNN) class TestRKNNModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - import mmdeploy.backend.rknn as rknn_apis - from mmdeploy.backend.rknn import RKNNWrapper - rknn_apis.__dict__.update({'RKNNWrapper': RKNNWrapper}) + from mmdeploy.backend.rknn.wrapper import RKNNWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(RKNNWrapper) - cls.outputs = [ - torch.rand(1, 255, 5, 5), - torch.rand(1, 255, 10, 10), - torch.rand(1, 255, 20, 20) - ] - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config({ - 'onnx_config': { - 'output_names': ['output'] - }, - 'backend_config': { - 'common_config': {} - } - }) - model_cfg = Config( - dict( - model=dict( - bbox_head=dict( - type='YOLOV3Head', - num_classes=80, - in_channels=[512, 256, 128], - out_channels=[1024, 512, 256], - anchor_generator=dict( - type='YOLOAnchorGenerator', - base_sizes=[[(116, 90), (156, 198), ( - 373, 326)], [(30, 61), (62, 45), ( - 59, 119)], [(10, 13), (16, 30), (33, 23)]], - strides=[32, 16, 8]), - bbox_coder=dict(type='YOLOBBoxCoder'), - featmap_strides=[32, 16, 8], - loss_cls=dict( - type='CrossEntropyLoss', - use_sigmoid=True, - loss_weight=1.0, - reduction='sum'), - loss_conf=dict( - type='CrossEntropyLoss', - use_sigmoid=True, - loss_weight=1.0, - reduction='sum'), - loss_xy=dict( - type='CrossEntropyLoss', - use_sigmoid=True, - loss_weight=2.0, - reduction='sum'), - loss_wh=dict( - type='MSELoss', loss_weight=2.0, reduction='sum')), - test_cfg=dict( - nms_pre=1000, - min_bbox_size=0, - score_thr=0.05, - conf_thr=0.005, - nms=dict(type='nms', iou_threshold=0.45), - max_per_img=100)))) - - from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ - RKNNModel - cls.rknn_model = RKNNModel(Backend.RKNN, ['', ''], 'cpu', - ['' for i in range(80)], model_cfg, - deploy_cfg) - - @classmethod - def teardown_class(cls): - cls.wrapper.recover() - - def test_forward_test(self): + with SwitchBackendWrapper(RKNNWrapper) as wrapper: + outputs = [ + torch.rand(1, 255, 5, 5), + torch.rand(1, 255, 10, 10), + torch.rand(1, 255, 20, 20) + ] + wrapper.set(outputs=outputs) + deploy_cfg = Config({ + 'onnx_config': { + 'output_names': ['output'] + }, + 'backend_config': { + 'common_config': {} + } + }) + model_cfg = Config( + dict( + model=dict( + bbox_head=dict( + type='YOLOV3Head', + num_classes=80, + in_channels=[512, 256, 128], + out_channels=[1024, 512, 256], + anchor_generator=dict( + type='YOLOAnchorGenerator', + base_sizes=[[(116, 90), (156, 198), ( + 373, 326)], [(30, 61), (62, 45), ( + 59, + 119)], [(10, 13), (16, 30), (33, 23)]], + strides=[32, 16, 8]), + bbox_coder=dict(type='YOLOBBoxCoder'), + featmap_strides=[32, 16, 8], + loss_cls=dict( + type='CrossEntropyLoss', + use_sigmoid=True, + loss_weight=1.0, + reduction='sum'), + loss_conf=dict( + type='CrossEntropyLoss', + use_sigmoid=True, + loss_weight=1.0, + reduction='sum'), + loss_xy=dict( + type='CrossEntropyLoss', + use_sigmoid=True, + loss_weight=2.0, + reduction='sum'), + loss_wh=dict( + type='MSELoss', + loss_weight=2.0, + reduction='sum')), + test_cfg=dict( + nms_pre=1000, + min_bbox_size=0, + score_thr=0.05, + conf_thr=0.005, + nms=dict(type='nms', iou_threshold=0.45), + max_per_img=100)))) + + from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ + RKNNModel + yield RKNNModel(Backend.RKNN, ['', ''], 'cpu', + ['' for i in range(80)], model_cfg, deploy_cfg) + + def test_forward_test(self, end2end_model): imgs = torch.rand(1, 3, 64, 64) - results = self.rknn_model.forward_test(imgs) + results = end2end_model.forward_test(imgs) assert_det_results(results, 'RKNNWrapper') diff --git a/tests/test_codebase/test_mmdet3d/test_voxel_detection.py b/tests/test_codebase/test_mmdet3d/test_voxel_detection.py index 5d60e81aa5..d09c72dc76 100644 --- a/tests/test_codebase/test_mmdet3d/test_voxel_detection.py +++ b/tests/test_codebase/test_mmdet3d/test_voxel_detection.py @@ -5,7 +5,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -51,19 +50,16 @@ def test_build_pytorch_model(): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set( - outputs={ - 'cls_score': torch.rand(1, 18, 32, 32), - 'bbox_pred': torch.rand(1, 42, 32, 32), - 'dir_cls_pred': torch.rand(1, 12, 32, 32) - }) - - yield task_processor.build_backend_model(['']) - - wrapper.recover() + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set( + outputs={ + 'cls_score': torch.rand(1, 18, 32, 32), + 'bbox_pred': torch.rand(1, 42, 32, 32), + 'dir_cls_pred': torch.rand(1, 12, 32, 32) + }) + + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmdet3d/test_voxel_detection_model.py b/tests/test_codebase/test_mmdet3d/test_voxel_detection_model.py index 5e651c4f2c..8cc0bd1c2b 100644 --- a/tests/test_codebase/test_mmdet3d/test_voxel_detection_model.py +++ b/tests/test_codebase/test_mmdet3d/test_voxel_detection_model.py @@ -3,7 +3,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -23,48 +22,43 @@ @backend_checker(Backend.ONNXRUNTIME) class TestVoxelDetectionModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'cls_score': torch.rand(1, 18, 32, 32), - 'bbox_pred': torch.rand(1, 42, 32, 32), - 'dir_cls_pred': torch.rand(1, 12, 32, 32) - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = mmengine.Config({ - 'onnx_config': { - 'input_names': ['voxels', 'num_points', 'coors'], - 'output_names': ['cls_score', 'bbox_pred', 'dir_cls_pred'], - 'opset_version': 11 - }, - 'backend_config': { - 'type': 'onnxruntime' + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'cls_score': torch.rand(1, 18, 32, 32), + 'bbox_pred': torch.rand(1, 42, 32, 32), + 'dir_cls_pred': torch.rand(1, 12, 32, 32) } - }) - - from mmdeploy.utils import load_config - model_cfg_path = 'tests/test_codebase/test_mmdet3d/data/model_cfg.py' - model_cfg = load_config(model_cfg_path)[0] - cls.end2end_model = VoxelDetectionModel( - Backend.ONNXRUNTIME, [''], - device='cuda', - deploy_cfg=deploy_cfg, - model_cfg=model_cfg) + wrapper.set(outputs=outputs) + deploy_cfg = mmengine.Config({ + 'onnx_config': { + 'input_names': ['voxels', 'num_points', 'coors'], + 'output_names': ['cls_score', 'bbox_pred', 'dir_cls_pred'], + 'opset_version': 11 + }, + 'backend_config': { + 'type': 'onnxruntime' + } + }) - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + from mmdeploy.utils import load_config + model_cfg_path = 'tests/test_codebase/test_mmdet3d/data/model_cfg.py' # noqa + model_cfg = load_config(model_cfg_path)[0] + yield VoxelDetectionModel( + Backend.ONNXRUNTIME, [''], + device='cuda', + deploy_cfg=deploy_cfg, + model_cfg=model_cfg) @pytest.mark.skipif( reason='Only support GPU test', condition=not torch.cuda.is_available()) - def test_forward_and_show_result(self): + def test_forward_and_show_result(self, end2end_model): inputs = { 'voxels': { 'voxels': torch.rand((3945, 32, 4)), @@ -72,7 +66,7 @@ def test_forward_and_show_result(self): 'coors': torch.ones((3945, 4), dtype=torch.int32) } } - results = self.end2end_model.forward(inputs=inputs) + results = end2end_model.forward(inputs=inputs) assert results is not None @@ -88,8 +82,7 @@ def test_build_voxel_detection_model(): output_names=['cls_score', 'bbox_pred', 'dir_cls_pred']), codebase_config=dict(type=Codebase.MMDET3D.value))) - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_codebase/test_mmedit/test_super_resolution.py b/tests/test_codebase/test_mmedit/test_super_resolution.py index d0a82e29b8..80591e7341 100644 --- a/tests/test_codebase/test_mmedit/test_super_resolution.py +++ b/tests/test_codebase/test_mmedit/test_super_resolution.py @@ -9,7 +9,6 @@ from torch.utils.data import DataLoader from torch.utils.data.dataset import Dataset -import mmdeploy.apis.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.core.rewriters.rewriter_manager import RewriterContext @@ -51,16 +50,13 @@ def init_task_processor(): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'output': torch.rand(1, 3, 50, 50), - }) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(outputs={ + 'output': torch.rand(1, 3, 50, 50), + }) - yield task_processor.build_backend_model(['']) - - wrapper.recover() + yield task_processor.build_backend_model(['']) def test_build_test_runner(): diff --git a/tests/test_codebase/test_mmedit/test_super_resolution_model.py b/tests/test_codebase/test_mmedit/test_super_resolution_model.py index 0c8a1dc52f..4ccbaa1c1d 100644 --- a/tests/test_codebase/test_mmedit/test_super_resolution_model.py +++ b/tests/test_codebase/test_mmedit/test_super_resolution_model.py @@ -21,7 +21,7 @@ class TestEnd2EndModel: def end2end_model(self): # force add backend wrapper regardless of plugins # make sure ONNXRuntimeEditor can use ORTWrapper inside itself - from mmdeploy.backend.onnxruntime import ORTWrapper + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper from mmdeploy.codebase.mmedit.deploy.super_resolution_model import \ End2EndModel diff --git a/tests/test_codebase/test_mmocr/test_text_detection.py b/tests/test_codebase/test_mmocr/test_text_detection.py index 57f2db7281..0d95d236b2 100644 --- a/tests/test_codebase/test_mmocr/test_text_detection.py +++ b/tests/test_codebase/test_mmocr/test_text_detection.py @@ -7,7 +7,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -56,16 +55,13 @@ def test_build_pytorch_model(): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'output': torch.rand(1, *img_shape), - }) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(outputs={ + 'output': torch.rand(1, *img_shape), + }) - yield task_processor.build_backend_model(['']) - - wrapper.recover() + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmocr/test_text_detection_models.py b/tests/test_codebase/test_mmocr/test_text_detection_models.py index ff32007472..03f0659024 100644 --- a/tests/test_codebase/test_mmocr/test_text_detection_models.py +++ b/tests/test_codebase/test_mmocr/test_text_detection_models.py @@ -3,7 +3,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase, load_config from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -19,41 +18,36 @@ @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'outputs': torch.rand(1, IMAGE_SIZE, IMAGE_SIZE), - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = mmengine.Config( - {'onnx_config': { - 'output_names': ['outputs'] - }}) - model_cfg_path = 'tests/test_codebase/test_mmocr/data/dbnet.py' - model_cfg = load_config(model_cfg_path)[0] - - from mmdeploy.codebase.mmocr.deploy.text_detection_model import \ - End2EndModel - cls.end2end_model = End2EndModel( - Backend.ONNXRUNTIME, [''], - device='cpu', - deploy_cfg=deploy_cfg, - model_cfg=model_cfg) + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'outputs': torch.rand(1, IMAGE_SIZE, IMAGE_SIZE), + } + wrapper.set(outputs=outputs) + deploy_cfg = mmengine.Config( + {'onnx_config': { + 'output_names': ['outputs'] + }}) + model_cfg_path = 'tests/test_codebase/test_mmocr/data/dbnet.py' + model_cfg = load_config(model_cfg_path)[0] - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + from mmdeploy.codebase.mmocr.deploy.text_detection_model import \ + End2EndModel + yield End2EndModel( + Backend.ONNXRUNTIME, [''], + device='cpu', + deploy_cfg=deploy_cfg, + model_cfg=model_cfg) @pytest.mark.parametrize( 'ori_shape', [[IMAGE_SIZE, IMAGE_SIZE], [2 * IMAGE_SIZE, 2 * IMAGE_SIZE]]) - def test_forward(self, ori_shape): + def test_forward(self, end2end_model, ori_shape): imgs = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE) img_meta = { 'ori_shape': ori_shape, @@ -66,7 +60,7 @@ def test_forward(self, ori_shape): pred_instances = InstanceData(metainfo=img_meta) data_sample = TextDetDataSample(pred_instances=pred_instances) data_sample.set_metainfo(img_meta) - results = self.end2end_model.forward(imgs, [data_sample]) + results = end2end_model.forward(imgs, [data_sample]) assert results is not None, 'failed to get output using '\ 'End2EndModel' @@ -81,8 +75,7 @@ def test_build_text_detection_model(): onnx_config=dict(output_names=['outputs']), codebase_config=dict(type='mmocr'))) - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_codebase/test_mmocr/test_text_recognition.py b/tests/test_codebase/test_mmocr/test_text_recognition.py index 8b81224098..8f5dd6fd87 100644 --- a/tests/test_codebase/test_mmocr/test_text_recognition.py +++ b/tests/test_codebase/test_mmocr/test_text_recognition.py @@ -7,7 +7,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -55,16 +54,13 @@ def test_build_pytorch_model(): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'output': torch.rand(1, 9, 37), - }) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(outputs={ + 'output': torch.rand(1, 9, 37), + }) - yield task_processor.build_backend_model(['']) - - wrapper.recover() + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmocr/test_text_recognition_models.py b/tests/test_codebase/test_mmocr/test_text_recognition_models.py index e2661e6f43..39e7000632 100644 --- a/tests/test_codebase/test_mmocr/test_text_recognition_models.py +++ b/tests/test_codebase/test_mmocr/test_text_recognition_models.py @@ -3,7 +3,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase, load_config from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -19,41 +18,36 @@ @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'output': torch.rand(1, 9, 37), - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = mmengine.Config( - {'onnx_config': { - 'output_names': ['output'] - }}) - model_cfg_path = 'tests/test_codebase/test_mmocr/data/crnn.py' - model_cfg = load_config(model_cfg_path)[0] - - from mmdeploy.codebase.mmocr.deploy.text_recognition_model import \ - End2EndModel - cls.end2end_model = End2EndModel( - Backend.ONNXRUNTIME, [''], - device='cpu', - deploy_cfg=deploy_cfg, - model_cfg=model_cfg) + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'output': torch.rand(1, 9, 37), + } + wrapper.set(outputs=outputs) + deploy_cfg = mmengine.Config( + {'onnx_config': { + 'output_names': ['output'] + }}) + model_cfg_path = 'tests/test_codebase/test_mmocr/data/crnn.py' + model_cfg = load_config(model_cfg_path)[0] - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + from mmdeploy.codebase.mmocr.deploy.text_recognition_model import \ + End2EndModel + yield End2EndModel( + Backend.ONNXRUNTIME, [''], + device='cpu', + deploy_cfg=deploy_cfg, + model_cfg=model_cfg) @pytest.mark.parametrize( 'ori_shape', [[IMAGE_SIZE, IMAGE_SIZE, 3], [2 * IMAGE_SIZE, 2 * IMAGE_SIZE, 3]]) - def test_forward(self, ori_shape): + def test_forward(self, end2end_model, ori_shape): imgs = [torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE)] img_meta = { 'ori_shape': ori_shape, @@ -65,7 +59,7 @@ def test_forward(self, ori_shape): pred_instances = InstanceData(metainfo=img_meta) data_sample = TextRecogDataSample(pred_instances=pred_instances) data_sample.set_metainfo(img_meta) - results = self.end2end_model.forward(imgs, [data_sample]) + results = end2end_model.forward(imgs, [data_sample]) assert results is not None, 'failed to get output using '\ 'End2EndModel' @@ -80,8 +74,7 @@ def test_build_text_recognition_model(): onnx_config=dict(output_names=['outputs']), codebase_config=dict(type='mmocr'))) - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_codebase/test_mmpose/test_pose_detection.py b/tests/test_codebase/test_mmpose/test_pose_detection.py index 1d9b12e365..9702675690 100644 --- a/tests/test_codebase/test_mmpose/test_pose_detection.py +++ b/tests/test_codebase/test_mmpose/test_pose_detection.py @@ -5,7 +5,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config from mmdeploy.utils.test import SwitchBackendWrapper @@ -48,16 +47,14 @@ def test_build_pytorch_model(): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'output': torch.rand(1, num_output_channels, *heatmap_shape), - }) - - yield task_processor.build_backend_model(['']) - - wrapper.recover() + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set( + outputs={ + 'output': torch.rand(1, num_output_channels, *heatmap_shape), + }) + + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmpose/test_pose_detection_model.py b/tests/test_codebase/test_mmpose/test_pose_detection_model.py index 89baca6832..7419388b8c 100644 --- a/tests/test_codebase/test_mmpose/test_pose_detection_model.py +++ b/tests/test_codebase/test_mmpose/test_pose_detection_model.py @@ -2,7 +2,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase, load_config from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -23,38 +22,33 @@ @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'output': torch.rand(1, 1, IMAGE_H, IMAGE_W), - } - cls.wrapper.set(outputs=cls.outputs) + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'output': torch.rand(1, 1, IMAGE_H, IMAGE_W), + } + wrapper.set(outputs=outputs) - from mmdeploy.codebase.mmpose.deploy.pose_detection_model import \ - End2EndModel - model_cfg_path = 'tests/test_codebase/test_mmpose/data/model.py' - model_cfg = load_config(model_cfg_path)[0] - deploy_cfg = generate_mmpose_deploy_config() - cls.end2end_model = End2EndModel( - Backend.ONNXRUNTIME, [''], - device='cpu', - deploy_cfg=deploy_cfg, - model_cfg=model_cfg) + from mmdeploy.codebase.mmpose.deploy.pose_detection_model import \ + End2EndModel + model_cfg_path = 'tests/test_codebase/test_mmpose/data/model.py' + model_cfg = load_config(model_cfg_path)[0] + deploy_cfg = generate_mmpose_deploy_config() + yield End2EndModel( + Backend.ONNXRUNTIME, [''], + device='cpu', + deploy_cfg=deploy_cfg, + model_cfg=model_cfg) - @classmethod - def teardown_class(cls): - cls.wrapper.recover() - - def test_forward(self): + def test_forward(self, end2end_model): img = torch.rand(1, 3, IMAGE_H, IMAGE_W) data_samples = [generate_datasample((IMAGE_H, IMAGE_W))] - results = self.end2end_model.forward(img, data_samples) + results = end2end_model.forward(img, data_samples) assert results is not None, 'failed to get output using '\ 'End2EndModel' @@ -65,8 +59,7 @@ def test_build_pose_detection_model(): model_cfg = load_config(model_cfg_path)[0] deploy_cfg = generate_mmpose_deploy_config() - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_codebase/test_mmrotate/test_rotated_detection.py b/tests/test_codebase/test_mmrotate/test_rotated_detection.py index 206857becf..bcfc86b267 100644 --- a/tests/test_codebase/test_mmrotate/test_rotated_detection.py +++ b/tests/test_codebase/test_mmrotate/test_rotated_detection.py @@ -9,7 +9,6 @@ from torch.utils.data import DataLoader from torch.utils.data.dataset import Dataset -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -62,17 +61,14 @@ def test_build_pytorch_model(): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'dets': torch.rand(1, 10, 6), - 'labels': torch.randint(1, 10, (1, 10)) - }) - - yield task_processor.build_backend_model(['']) - - wrapper.recover() + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(outputs={ + 'dets': torch.rand(1, 10, 6), + 'labels': torch.randint(1, 10, (1, 10)) + }) + + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmrotate/test_rotated_detection_model.py b/tests/test_codebase/test_mmrotate/test_rotated_detection_model.py index 9aa710bc2f..0b026db57b 100644 --- a/tests/test_codebase/test_mmrotate/test_rotated_detection_model.py +++ b/tests/test_codebase/test_mmrotate/test_rotated_detection_model.py @@ -4,7 +4,6 @@ from mmengine import Config from mmengine.structures import BaseDataElement -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase, load_config from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -21,34 +20,29 @@ @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - - # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'dets': torch.rand(1, 10, 6), - 'labels': torch.rand(1, 10) - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = Config( - {'onnx_config': { - 'output_names': ['dets', 'labels'] - }}) - + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper from mmdeploy.codebase.mmrotate.deploy.rotated_detection_model import \ End2EndModel - cls.end2end_model = End2EndModel( - Backend.ONNXRUNTIME, [''], device='cpu', deploy_cfg=deploy_cfg) - @classmethod - def teardown_class(cls): - cls.wrapper.recover() + # simplify backend inference + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'dets': torch.rand(1, 10, 6), + 'labels': torch.rand(1, 10) + } + wrapper.set(outputs=outputs) + deploy_cfg = Config( + {'onnx_config': { + 'output_names': ['dets', 'labels'] + }}) + + yield End2EndModel( + Backend.ONNXRUNTIME, [''], device='cpu', deploy_cfg=deploy_cfg) - def test_forward(self): + def test_forward(self, end2end_model): imgs = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE) img_metas = [ BaseDataElement(metainfo={ @@ -56,7 +50,7 @@ def test_forward(self): 'scale_factor': [1, 1] }) ] - results = self.end2end_model.forward(imgs, img_metas) + results = end2end_model.forward(imgs, img_metas) assert results is not None, 'failed to get output using End2EndModel' @@ -70,8 +64,7 @@ def test_build_rotated_detection_model(): ir_config=dict(type='onnx', output_names=['dets', 'labels']), codebase_config=dict(type='mmrotate'))) - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_codebase/test_mmseg/test_segmentation.py b/tests/test_codebase/test_mmseg/test_segmentation.py index 0ecd88c374..0cd924b28e 100644 --- a/tests/test_codebase/test_mmseg/test_segmentation.py +++ b/tests/test_codebase/test_mmseg/test_segmentation.py @@ -7,7 +7,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.apis import build_task_processor from mmdeploy.codebase import import_codebase from mmdeploy.utils import Codebase, load_config @@ -70,16 +69,13 @@ def test_build_pytorch_model(from_mmrazor: Any): @pytest.fixture def backend_model(): - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - wrapper = SwitchBackendWrapper(ORTWrapper) - wrapper.set(outputs={ - 'output': torch.rand(1, 1, *img_shape), - }) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(outputs={ + 'output': torch.rand(1, 1, *img_shape), + }) - yield task_processor.build_backend_model(['']) - - wrapper.recover() + yield task_processor.build_backend_model(['']) def test_build_backend_model(backend_model): diff --git a/tests/test_codebase/test_mmseg/test_segmentation_model.py b/tests/test_codebase/test_mmseg/test_segmentation_model.py index 5ba41d2a30..5bb81c537a 100644 --- a/tests/test_codebase/test_mmseg/test_segmentation_model.py +++ b/tests/test_codebase/test_mmseg/test_segmentation_model.py @@ -4,7 +4,6 @@ import pytest import torch -import mmdeploy.backend.onnxruntime as ort_apis from mmdeploy.codebase import import_codebase from mmdeploy.utils import Backend, Codebase from mmdeploy.utils.test import SwitchBackendWrapper, backend_checker @@ -24,34 +23,29 @@ @backend_checker(Backend.ONNXRUNTIME) class TestEnd2EndModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(ORTWrapper) - cls.outputs = { - 'output': torch.rand(1, 1, IMAGE_SIZE, IMAGE_SIZE), - } - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = generate_mmseg_deploy_config() - - from mmdeploy.codebase.mmseg.deploy.segmentation_model import \ - End2EndModel - cls.end2end_model = End2EndModel( - Backend.ONNXRUNTIME, [''], device='cpu', deploy_cfg=deploy_cfg) - - @classmethod - def teardown_class(cls): - cls.wrapper.recover() - - def test_forward(self): + with SwitchBackendWrapper(ORTWrapper) as wrapper: + outputs = { + 'output': torch.rand(1, 1, IMAGE_SIZE, IMAGE_SIZE), + } + wrapper.set(outputs=outputs) + deploy_cfg = generate_mmseg_deploy_config() + + from mmdeploy.codebase.mmseg.deploy.segmentation_model import \ + End2EndModel + yield End2EndModel( + Backend.ONNXRUNTIME, [''], device='cpu', deploy_cfg=deploy_cfg) + + def test_forward(self, end2end_model): from mmseg.structures import SegDataSample imgs = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE) data_samples = [generate_datasample(IMAGE_SIZE, IMAGE_SIZE)] - results = self.end2end_model.forward(imgs, data_samples) + results = end2end_model.forward(imgs, data_samples) assert len(results) == 1 assert isinstance(results[0], SegDataSample) @@ -59,39 +53,38 @@ def test_forward(self): @backend_checker(Backend.RKNN) class TestRKNNModel: - @classmethod - def setup_class(cls): + @pytest.fixture(scope='class') + def end2end_model(self): # force add backend wrapper regardless of plugins - import mmdeploy.backend.rknn as rknn_apis - from mmdeploy.backend.rknn import RKNNWrapper - rknn_apis.__dict__.update({'RKNNWrapper': RKNNWrapper}) + from mmdeploy.backend.rknn.wrapper import RKNNWrapper # simplify backend inference - cls.wrapper = SwitchBackendWrapper(RKNNWrapper) - cls.outputs = [torch.rand(1, 19, IMAGE_SIZE, IMAGE_SIZE)] - cls.wrapper.set(outputs=cls.outputs) - deploy_cfg = mmengine.Config({ - 'onnx_config': { - 'output_names': ['outputs'] - }, - 'backend_config': { - 'common_config': {} - } - }) - - from mmdeploy.codebase.mmseg.deploy.segmentation_model import RKNNModel - class_names = ['' for i in range(NUM_CLASS)] - palette = np.random.randint(0, 255, size=(NUM_CLASS, 3)) - cls.rknn_model = RKNNModel( - Backend.RKNN, [''], - device='cpu', - class_names=class_names, - palette=palette, - deploy_cfg=deploy_cfg) - - def test_forward_test(self): + with SwitchBackendWrapper(RKNNWrapper) as wrapper: + outputs = [torch.rand(1, 19, IMAGE_SIZE, IMAGE_SIZE)] + wrapper.set(outputs=outputs) + deploy_cfg = mmengine.Config({ + 'onnx_config': { + 'output_names': ['outputs'] + }, + 'backend_config': { + 'common_config': {} + } + }) + + from mmdeploy.codebase.mmseg.deploy.segmentation_model import \ + RKNNModel + class_names = ['' for i in range(NUM_CLASS)] + palette = np.random.randint(0, 255, size=(NUM_CLASS, 3)) + yield RKNNModel( + Backend.RKNN, [''], + device='cpu', + class_names=class_names, + palette=palette, + deploy_cfg=deploy_cfg) + + def test_forward_test(self, end2end_model): imgs = torch.rand(2, 3, IMAGE_SIZE, IMAGE_SIZE) - results = self.rknn_model.forward_test(imgs) + results = end2end_model.forward_test(imgs) assert isinstance(results[0], np.ndarray) @@ -101,8 +94,7 @@ def test_build_segmentation_model(): dict(data=dict(test={'type': 'CityscapesDataset'}))) deploy_cfg = generate_mmseg_deploy_config() - from mmdeploy.backend.onnxruntime import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper # simplify backend inference with SwitchBackendWrapper(ORTWrapper) as wrapper: diff --git a/tests/test_ir/test_onnx.py b/tests/test_ir/test_onnx.py new file mode 100644 index 0000000000..f81cdebfc2 --- /dev/null +++ b/tests/test_ir/test_onnx.py @@ -0,0 +1,78 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import importlib +import os.path as osp + +import pytest + +from mmdeploy.ir.onnx import ONNXManager, ONNXParam + + +@pytest.fixture(scope='class') +def file_name(): + return 'tmp' + + +@pytest.fixture(scope='class') +def dummy_model(): + import torch + + class DummyModel(torch.nn.Module): + + def forward(self, x, y): + return x + y, x - y + + return DummyModel() + + +@pytest.fixture(scope='class') +def dummy_args(): + import torch + return (torch.ones([2, 2]), torch.ones([2, 2]) + 1) + + +@pytest.fixture +def dummy_param(dummy_args, tmp_path, file_name): + return ONNXParam( + args=dummy_args, + work_dir=str(tmp_path), + file_name=file_name, + input_names=['x', 'y'], + output_names=['out0', 'out1']) + + +@pytest.mark.skipif( + importlib.util.find_spec('torch') is None, reason='PyTorch is required.') +class TestONNXParam: + + def test_file_name(self, dummy_param, file_name): + assert dummy_param.file_name == f'{file_name}.onnx' + + def test_check(self, dummy_param): + dummy_param.check() + + with pytest.raises(AssertionError): + dummy_param.opset_version = 6 + dummy_param.check() + + +@pytest.mark.skipif( + importlib.util.find_spec('torch') is None, reason='PyTorch is required.') +class TestONNXManager: + + def test_build_param(self, dummy_args, file_name): + assert isinstance( + ONNXManager.build_param(args=dummy_args, file_name=file_name), + ONNXParam) + + def test_is_available(self): + assert ONNXManager.is_available() + + def test_export(self, dummy_model, dummy_args, file_name, tmp_path): + output_path = str(tmp_path / f'{file_name}.onnx') + ONNXManager.export(dummy_model, dummy_args, output_path) + assert osp.exists(output_path) + + def test_export_from_param(self, dummy_model, dummy_param): + output_path = osp.join(dummy_param.work_dir, dummy_param.file_name) + ONNXManager.export_from_param(dummy_model, dummy_param) + assert osp.exists(output_path) diff --git a/tests/test_ir/test_torchscript.py b/tests/test_ir/test_torchscript.py new file mode 100644 index 0000000000..2e4e728a6d --- /dev/null +++ b/tests/test_ir/test_torchscript.py @@ -0,0 +1,67 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import importlib +import os.path as osp + +import pytest + +from mmdeploy.ir.torchscript import TorchScriptManager, TorchScriptParam + + +@pytest.fixture(scope='class') +def file_name(): + return 'tmp' + + +@pytest.fixture(scope='class') +def dummy_model(): + import torch + + class DummyModel(torch.nn.Module): + + def forward(self, x, y): + return x + y, x - y + + return DummyModel() + + +@pytest.fixture(scope='class') +def dummy_args(): + import torch + return (torch.ones([2, 2]), torch.ones([2, 2]) + 1) + + +@pytest.fixture +def dummy_param(dummy_args, tmp_path, file_name): + return TorchScriptParam( + args=dummy_args, work_dir=str(tmp_path), file_name=file_name) + + +@pytest.mark.skipif( + importlib.util.find_spec('torch') is None, reason='PyTorch is required.') +class TestTorchScriptParam: + + def test_file_name(self, dummy_param, file_name): + assert dummy_param.file_name == f'{file_name}.pth' + + +@pytest.mark.skipif( + importlib.util.find_spec('torch') is None, reason='PyTorch is required.') +class TestTorchScriptManager: + + def test_build_param(self, dummy_args, file_name): + assert isinstance( + TorchScriptManager.build_param( + args=dummy_args, file_name=file_name), TorchScriptParam) + + def test_is_available(self): + assert TorchScriptManager.is_available() + + def test_export(self, dummy_model, dummy_args, file_name, tmp_path): + output_path = str(tmp_path / f'{file_name}.pth') + TorchScriptManager.export(dummy_model, dummy_args, output_path) + assert osp.exists(output_path) + + def test_export_from_param(self, dummy_model, dummy_param): + output_path = osp.join(dummy_param.work_dir, dummy_param.file_name) + TorchScriptManager.export_from_param(dummy_model, dummy_param) + assert osp.exists(output_path) diff --git a/tests/test_ops/test_ops.py b/tests/test_ops/test_ops.py index 6cee9acb3b..89c1cb3c80 100644 --- a/tests/test_ops/test_ops.py +++ b/tests/test_ops/test_ops.py @@ -769,7 +769,7 @@ def test_gather(backend, assert importlib.util.find_spec('onnxruntime') is not None, 'onnxruntime \ not installed.' - from mmdeploy.backend.onnxruntime import ORTWrapper + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper ort_model = ORTWrapper( gather_model.SerializeToString(), device='cpu', @@ -1212,8 +1212,8 @@ def test_multiclass_nms_rotated_with_keep_top_k(backend, pre_top_k): test_scores = torch.ones(batch_size, num_boxes, num_classes) model_inputs = {'boxes': test_boxes, 'scores': test_scores} - import mmdeploy.backend.onnxruntime as ort_apis - backend_model = ort_apis.ORTWrapper(onnx_model_path, 'cpu', None) + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper + backend_model = ORTWrapper(onnx_model_path, 'cpu', None) output = backend_model.forward(model_inputs) output = backend_model.output_to_list(output) dets = output[0] diff --git a/tests/test_ops/utils.py b/tests/test_ops/utils.py index c36f678e09..390237a0d8 100644 --- a/tests/test_ops/utils.py +++ b/tests/test_ops/utils.py @@ -61,7 +61,7 @@ def run_and_validate(self, else: model_outputs = list(model_outputs) - from mmdeploy.backend.onnxruntime import ORTWrapper + from mmdeploy.backend.onnxruntime.wrapper import ORTWrapper onnx_model = ORTWrapper(onnx_file_path, 'cpu', output_names) with torch.no_grad(): onnx_outputs = onnx_model.forward( @@ -151,7 +151,7 @@ def run_and_validate(self, else: model_outputs = [data.cpu().float() for data in model_outputs] - from mmdeploy.backend.tensorrt import TRTWrapper + from mmdeploy.backend.tensorrt.wrapper import TRTWrapper trt_model = TRTWrapper(trt_file_path, output_names) trt_outputs = trt_model(dict(zip(input_names, input_list))) trt_outputs = [trt_outputs[i].float().cpu() for i in output_names] @@ -214,7 +214,7 @@ def run_and_validate(self, model_output.float() for model_output in model_outputs ] - from mmdeploy.backend.ncnn import NCNNWrapper + from mmdeploy.backend.ncnn.wrapper import NCNNWrapper ncnn_model = NCNNWrapper(ncnn_param_path, ncnn_bin_path, output_names) ncnn_outputs = ncnn_model(dict(zip(input_names, inputs_list))) ncnn_outputs = [ncnn_outputs[name] for name in output_names] @@ -241,9 +241,9 @@ def _from_onnx(self, model, model_name, output_names, save_dir=None): onnx.save_model(model, onnx_file_path) from mmdeploy.backend.ncnn import from_onnx - from_onnx(onnx_file_path, os.path.join(save_dir, model_name)) + from_onnx(onnx_file_path, ncnn_param_path, ncnn_bin_path) - from mmdeploy.backend.ncnn import NCNNWrapper + from mmdeploy.backend.ncnn.wrapper import NCNNWrapper ncnn_model = NCNNWrapper(ncnn_param_path, ncnn_bin_path, output_names) diff --git a/tests/test_pytorch/test_pytorch_functions.py b/tests/test_pytorch/test_pytorch_functions.py index 245bfba9d4..4d8686f6da 100644 --- a/tests/test_pytorch/test_pytorch_functions.py +++ b/tests/test_pytorch/test_pytorch_functions.py @@ -319,7 +319,7 @@ def triu_caller(*arg, **kwargs): [torch.rand(1, 16, 16), torch.rand(1, 3, 16, 16)]) @pytest.mark.parametrize('dim', [1, 2]) def test_normalize_ncnn(input, dim): - import mmdeploy.apis.ncnn as ncnn_apis + import mmdeploy.backend.ncnn as ncnn_apis from mmdeploy.utils.test import get_onnx_model def norm_func(input, dim): @@ -329,9 +329,9 @@ def norm_func(input, dim): model_inputs = {'input': input} ir_file_path = get_onnx_model(wrapped_func, model_inputs, deploy_cfg_ncnn) assert osp.exists(ir_file_path) - ncnn_files_prefix = osp.splitext(ir_file_path)[0] - ncnn_apis.from_onnx(ir_file_path, ncnn_files_prefix) - param_path, bin_path = ncnn_apis.get_output_model_file(ir_file_path) + param_path, bin_path = ncnn_apis.onnx2ncnn.get_output_model_file( + ir_file_path) + ncnn_apis.to_backend(ir_file_path, param_path, bin_path) assert osp.exists(param_path) assert osp.exists(bin_path) diff --git a/tests/test_utils/test_docstring_parser.py b/tests/test_utils/test_docstring_parser.py new file mode 100644 index 0000000000..0263de92ec --- /dev/null +++ b/tests/test_utils/test_docstring_parser.py @@ -0,0 +1,194 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest + +from mmdeploy.utils import docstring_parser as parser + + +@pytest.fixture(scope='module') +def singleline_arg(): + return ' input_names (list, optional): input_names. Defaults to None.' + + +@pytest.fixture(scope='module') +def multiline_arg(): + return ( + ' input_names (list, optional): input_names. Defaults to None.\n' + ' next line.\n' + ' next next line.\n') + + +def test_parse_arg(singleline_arg, multiline_arg): + doc_arg, doc_len = parser.parse_arg(singleline_arg) + assert doc_arg is not None + assert doc_arg.name == 'input_names' + assert doc_arg.type == 'list, optional' + assert doc_arg.desc == 'input_names. Defaults to None.' + assert doc_len == len(singleline_arg) + + doc_arg, doc_len = parser.parse_arg(multiline_arg) + assert doc_arg is not None + desc = 'input_names. Defaults to None.next line.next next line.' + assert doc_arg.name == 'input_names' + assert doc_arg.type == 'list, optional' + assert doc_arg.desc == desc + assert doc_len == len(multiline_arg) + + +@pytest.fixture(scope='module') +def doc_args(): + return ( + ' args (Any): args\n' + ' input_names (list, optional): input_names. Defaults to None.\n' + ' next line.\n' + ' next next line.\n' + ' output_names (list, optional): output_names. Defaults to None.\n' + ' dynamic_axes (dict, optional): dynamic_axes. Defaults to None.\n' + ' backend (str, optional): backend. Defaults to `onnxruntime`.\n' + '\n\n' + 'Returns:\n' + ' int: return val') + + +def test_parse_args(doc_args): + doc_arg_list, doc_len = parser.parse_args(doc_args) + assert len(doc_arg_list) == 5 + assert doc_len == len(doc_args) - len('\n\n' + 'Returns:\n' + ' int: return val') + + gt_list = [ + ('args', 'Any', 'args'), + ('input_names', 'list, optional', + 'input_names. Defaults to None.next line.next next line.'), + ('output_names', 'list, optional', 'output_names. Defaults to None.'), + ('dynamic_axes', 'dict, optional', 'dynamic_axes. Defaults to None.'), + ('backend', 'str, optional', 'backend. Defaults to `onnxruntime`.'), + ] + + for doc_arg, gt in zip(doc_arg_list, gt_list): + gt_name, gt_type, gt_desc = gt + assert doc_arg.name == gt_name + assert doc_arg.type == gt_type + assert doc_arg.desc == gt_desc + + +@pytest.fixture(scope='module') +def empty_lines(): + return (' \n \nnot empty') + + +def test_parse_empty_line(empty_lines): + empty_len = parser.parse_empty_line(empty_lines) + assert empty_len == len(' \n \n') + assert empty_lines[empty_len:] == 'not empty' + + +@pytest.fixture(scope='module') +def doc_args_section(): + return ( + 'Args: \n' + ' \n \n' + ' args (Any): args\n' + ' input_names (list, optional): input_names. Defaults to None.\n' + ' next line.\n' + ' next next line.\n' + ' output_names (list, optional): output_names. Defaults to None.\n' + ' dynamic_axes (dict, optional): dynamic_axes. Defaults to None.\n' + ' backend (str, optional): backend. Defaults to `onnxruntime`.\n' + ' \n \n') + + +def test_args_section(doc_args_section): + doc_arg_list, doc_len = parser.parse_args_section(doc_args_section) + assert len(doc_arg_list) == 5 + assert doc_len == len(doc_args_section) + + gt_list = [ + ('args', 'Any', 'args'), + ('input_names', 'list, optional', + 'input_names. Defaults to None.next line.next next line.'), + ('output_names', 'list, optional', 'output_names. Defaults to None.'), + ('dynamic_axes', 'dict, optional', 'dynamic_axes. Defaults to None.'), + ('backend', 'str, optional', 'backend. Defaults to `onnxruntime`.'), + ] + + for doc_arg, gt in zip(doc_arg_list, gt_list): + gt_name, gt_type, gt_desc = gt + assert doc_arg.name == gt_name + assert doc_arg.type == gt_type + assert doc_arg.desc == gt_desc + + +@pytest.fixture(scope='module') +def full_doc_str(): + return ( + 'This is Head\n' + '\n' + 'This is desc1\n' + 'This is desc2\n' + '\n' + 'Args: \n' + ' \n \n' + ' args (Any): args\n' + ' input_names (list, optional): input_names. Defaults to None.\n' + ' next line.\n' + ' next next line.\n' + ' output_names (list, optional): output_names. Defaults to None.\n' + ' dynamic_axes (dict, optional): dynamic_axes. Defaults to None.\n' + ' backend (str, optional): backend. Defaults to `onnxruntime`.\n' + ' \n \n' + 'Returns:\n' + ' int: return val') + + +def test_parse_docstring(full_doc_str): + doc_str = parser.parse_docstring(full_doc_str) + + assert doc_str.head == 'This is Head' + assert doc_str.desc == 'This is desc1\nThis is desc2' + doc_arg_list = doc_str.args + + gt_list = [ + ('args', 'Any', 'args'), + ('input_names', 'list, optional', + 'input_names. Defaults to None.next line.next next line.'), + ('output_names', 'list, optional', 'output_names. Defaults to None.'), + ('dynamic_axes', 'dict, optional', 'dynamic_axes. Defaults to None.'), + ('backend', 'str, optional', 'backend. Defaults to `onnxruntime`.'), + ] + + for doc_arg, gt in zip(doc_arg_list, gt_list): + gt_name, gt_type, gt_desc = gt + assert doc_arg.name == gt_name + assert doc_arg.type == gt_type + assert doc_arg.desc == gt_desc + + +class TestInspectDocstringArguments: + + @pytest.fixture(scope='class') + def valid_obj(self): + + class ValidObj: + """Valid Obj. + + Args: + arg0 (int): description of arg0. + arg1 (bool): description of arg1. + """ + + return ValidObj + + def test_inspect(self, valid_obj): + assert parser.inspect_docstring_arguments(valid_obj) == [ + parser.DocStrArg( + name='arg0', type='int', desc='description of arg0.'), + parser.DocStrArg( + name='arg1', type='bool', desc='description of arg1.') + ] + + assert parser.inspect_docstring_arguments( + valid_obj, ignore_args=['arg0']) == [ + parser.DocStrArg( + name='arg1', type='bool', desc='description of arg1.') + ] diff --git a/tools/onnx2tensorrt.py b/tools/onnx2tensorrt.py index b7e7a7e505..b0e0bdad2c 100644 --- a/tools/onnx2tensorrt.py +++ b/tools/onnx2tensorrt.py @@ -2,8 +2,7 @@ import argparse import logging -from mmdeploy.backend.tensorrt import from_onnx -from mmdeploy.backend.tensorrt.utils import get_trt_log_level +from mmdeploy.backend.tensorrt.utils import from_onnx, get_trt_log_level from mmdeploy.utils import (get_common_config, get_model_inputs, get_root_logger, load_config)