diff --git a/paddlehub/compat/module/module_v1.py b/paddlehub/compat/module/module_v1.py index e93f13d313b435376c807b17b131e10f3909cd90..1ca899a13a89d04de499bd9571677037e653e166 100644 --- a/paddlehub/compat/module/module_v1.py +++ b/paddlehub/compat/module/module_v1.py @@ -18,6 +18,7 @@ import os from typing import Tuple, List import paddle +import paddle2onnx from easydict import EasyDict from paddlehub.compat import paddle_utils @@ -280,3 +281,30 @@ class ModuleV1(object): target_vars=list(fetch_dict.values()), model_filename=model_filename, params_filename=params_filename) + + def export_onnx_model(self, dirname: str, **kwargs): + ''' + Export the model to ONNX format. + + Args: + dirname(str): The directory to save the onnx model. + **kwargs(dict|optional): Other export configuration options for compatibility, some may be removed + in the future. Don't use them If not necessary. Refer to https://github.com/PaddlePaddle/paddle2onnx + for more information. + ''' + feed_dict, fetch_dict, program = self.context(for_test=True, trainable=False) + inputs = set([var.name for var in feed_dict.values()]) + if self.type == 'CV/classification': + outputs = [fetch_dict['class_probs']] + else: + outputs = set([var.name for var in fetch_dict.values()]) + outputs = [program.global_block().vars[key] for key in outputs] + + save_file = os.path.join(dirname, '{}.onnx'.format(self.name)) + paddle2onnx.program2onnx( + program=program, + scope=paddle.static.global_scope(), + feed_var_names=inputs, + target_vars=outputs, + save_file=save_file, + **kwargs) diff --git a/paddlehub/module/module.py b/paddlehub/module/module.py index d6eb377ff7587a718c24632d0be896c1b8180eda..38af30735dbfa193957458a3c43fd23a1fede219 100644 --- a/paddlehub/module/module.py +++ b/paddlehub/module/module.py @@ -22,9 +22,10 @@ import re import sys from typing import Callable, Generic, List, Optional, Union +import paddle +import paddle2onnx from easydict import EasyDict -import paddle from paddlehub.utils import parser, log, utils from paddlehub.compat import paddle_utils from paddlehub.compat.module.module_v1 import ModuleV1 @@ -131,6 +132,77 @@ class RunModule(object): def serving_func_name(self): return self._get_func_name(self.__class__, _module_serving_func) + @property + def _pretrained_model_path(self): + _pretrained_model_attrs = [ + 'pretrained_model_path', 'rec_pretrained_model_path', 'default_pretrained_model_path' + ] + + for _attr in _pretrained_model_attrs: + if hasattr(self, _attr): + path = getattr(self, _attr) + if os.path.exists(path) and os.path.isfile(path): + path = os.path.dirname(path) + return path + + return None + + def export_onnx_model(self, dirname: str, **kwargs): + ''' + Export the model to ONNX format. + + Args: + dirname(str): The directory to save the onnx model. + **kwargs(dict|optional): Other export configuration options for compatibility, some may be removed + in the future. Don't use them If not necessary. Refer to https://github.com/PaddlePaddle/paddle2onnx + for more information. + ''' + if not self._pretrained_model_path: + if isinstance(self, paddle.nn.Layer): + save_file = os.path.join(dirname, '{}'.format(self.name)) + if hasattr(self, 'input_spec'): + input_spec = self.input_sepc + else: + _type = self.type.lower() + if _type.startswith('cv/image'): + input_spec = paddle.static.InputSpec(shape=[None, 3, None, None], dtype='float32') + else: + raise NotImplementedError + paddle.onnx.export(self, save_file, input_spec=[input_spec]) + return + raise NotImplementedError + + place = paddle.CPUPlace() + exe = paddle.static.Executor(place) + + model_filename = None + params_filename = None + + if os.path.exists(os.path.join(self._pretrained_model_path, 'model')): + model_filename = 'model' + + if os.path.exists(os.path.join(self._pretrained_model_path, 'params')): + params_filename = 'params' + + if os.path.exists(os.path.join(self._pretrained_model_path, '__params__')): + params_filename = '__params__' + + save_file = os.path.join(dirname, '{}.onnx'.format(self.name)) + + program, inputs, outputs = paddle.fluid.io.load_inference_model( + dirname=self._pretrained_model_path, + model_filename=model_filename, + params_filename=params_filename, + executor=exe) + + paddle2onnx.program2onnx( + program=program, + scope=paddle.static.global_scope(), + feed_var_names=inputs, + target_vars=outputs, + save_file=save_file, + **kwargs) + class Module(object): ''' diff --git a/requirements.txt b/requirements.txt index 2721ee4b24d3d99e5038136e6e54abfcd56aca30..ead4c10b3de32873a77dbf7f1deec17754725bca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,8 @@ numpy matplotlib opencv-python packaging +paddle2onnx >= 0.5 +paddlenlp >= 2.0.0rc5 Pillow pyyaml pyzmq @@ -16,4 +18,3 @@ tqdm visualdl >= 2.0.0 # gunicorn not support windows gunicorn >= 19.10.0; sys_platform != "win32" -paddlenlp >= 2.0.0rc5