# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License" # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import import inspect import six import numpy as np import paddle from paddle import fluid from paddle.fluid import layers from paddle2onnx.legacy.graph import graph_helper, PaddleGraph from paddle2onnx.utils import logging from paddle2onnx.legacy.constant.op_mapping_status import * REGISTER_CUSTOM_PADDLE_OP = {} def get_max_support_version(versions, opset_version): max_version = -1 for vs in sorted(versions): if vs <= opset_version: max_version = vs return max_version def register_op_mapper(paddle_op, mapper_obj): paddle_op_list = [] if isinstance(paddle_op, six.string_types): paddle_op_list.append(paddle_op) elif isinstance(paddle_op, list): paddle_op_list = paddle_op else: raise ValueError('paddle_op must be List or string, but got type {}.'. format(type(paddle_op))) if not isinstance(mapper_obj, six.class_types): raise ValueError('mapper_obj must be Class, but got type {}.'.format( type(mapper_obj))) valid_register_func = 0 for k, v in inspect.getmembers(mapper_obj, inspect.ismethod): if k.startswith("opset_"): version = int(k.replace("opset_", "")) if version > 13 or version < 1: raise Exception( 'the specific method of operator mapper must be named opset_[number](1<=number<=13), such as opset_9, but got {}.'. format(k)) valid_register_func += 1 if valid_register_func == 0: raise Exception( 'the specific method of operator mapper must be classmethod, which named opset_[number](1<=number<=13), such as opset_9, but none achieved.' ) mapper = OpMapper(paddle_op_list) mapper(mapper_obj) class OpMapper(object): OPSETS = {} REGISTER_CUSTOM_PADDLE_OP = {} def __init__(self, paddle_op, **kwargs): if not isinstance(paddle_op, list): paddle_op = [paddle_op] self.paddle_op = paddle_op self.kwargs = kwargs def __call__(self, cls): for k, v in inspect.getmembers(cls, inspect.ismethod): if k.startswith("opset_"): version = int(k.replace("opset_", "")) for op in self.paddle_op: if op not in OpMapper.OPSETS: OpMapper.OPSETS[op] = {} opset_dict = OpMapper.OPSETS[op] opset_dict[version] = (v, self.kwargs) @staticmethod def mapping(graph, node, operator_export_type="ONNX"): try: if node.type in OpMapper.REGISTER_CUSTOM_PADDLE_OP: if operator_export_type in ["PaddleFallback"]: opsets = OpMapper.OPSETS[node.type] versions = list(opsets.keys()) convert_version = get_max_support_version( versions, graph.opset_version) mapper_func, kw = opsets[convert_version] mapper_func(graph, node, **kw) else: custom_paddle_op = OpMapper.REGISTER_CUSTOM_PADDLE_OP[ node.type](node) custom_paddle_graph, output_results = custom_paddle_op.get_paddle_graph( ) OpMapper.check_support_status(custom_paddle_graph.node_map, graph.opset_version) graph.build_op_nodes(custom_paddle_graph.node_map) node_output_results = dict() for k in node.output_names: custom_outs = output_results[k] node_outs = node.output(k) assert len(custom_outs) == len( node_outs ), "Length of custom implementation operator's outputs is not same with the length of original operator's outputs." for i in range(len(custom_outs)): graph.make_node( "Identity", inputs=[custom_outs[i]], outputs=[node_outs[i]]) else: opsets = OpMapper.OPSETS[node.type] versions = list(opsets.keys()) convert_version = get_max_support_version(versions, graph.opset_version) mapper_func, kw = opsets[convert_version] mapper_func(graph, node, **kw) except Exception as e: raise Exception( "Error happened when mapping node ['{}'] to onnx, which op_type is '{}' with inputs: {} and outputs: {}, specific error: ". format(node.layer_name, node.type, node.inputs, node.outputs) + str(e)) @staticmethod def get_recommend_opset_version(node_map, opset_version): recommend_opset_version = OpMapper.check_support_status( node_map, opset_version, True) for name, node in list(node_map.items()): if node.type in OpMapper.REGISTER_CUSTOM_PADDLE_OP: #如果是custom的op，获取custom的推荐op custom_paddle_op = OpMapper.REGISTER_CUSTOM_PADDLE_OP[ node.type](node) custom_paddle_graph, output_results = custom_paddle_op.get_paddle_graph( ) custom_recommend_opset_version = OpMapper.check_support_status( custom_paddle_graph.node_map, opset_version, True) recommend_opset_version = max(recommend_opset_version, custom_recommend_opset_version) if opset_version != recommend_opset_version: warning_info = "\n======================\n" warning_info += "\nFor a successful conversion, set the recommended opset version : {}\n".format( recommend_opset_version) warning_info += "\n======================\n" logging.warning(warning_info) return recommend_opset_version @staticmethod def check_support_status(node_map, opset_version, for_check=False): op_mapping_status = { OP_MAPPING_NO_REGISTER: [], OP_MAPPING_NO_VERSION: [], } for name, node in list(node_map.items()): if node.type in OpMapper.REGISTER_CUSTOM_PADDLE_OP: continue if node.type not in OpMapper.OPSETS: op_mapping_status[OP_MAPPING_NO_REGISTER].append(node) else: opsets = OpMapper.OPSETS[node.type] versions = list(opsets.keys()) convert_version = get_max_support_version(versions, opset_version) if convert_version == -1: op_mapping_status[OP_MAPPING_NO_VERSION].append(node) if len(op_mapping_status[OP_MAPPING_NO_REGISTER]) > 0: unsupported_op_types = set([ node.type for node in op_mapping_status[OP_MAPPING_NO_REGISTER] ]) error_info = "\nThere's {} ops are not supported yet\n".format( len(unsupported_op_types)) for op_type in unsupported_op_types: error_info += "=========== {} ===========\n".format(op_type) raise NotImplementedError(error_info) if len(op_mapping_status[OP_MAPPING_NO_VERSION]) > 0: unsupported_op_types = set([ node.type for node in op_mapping_status[OP_MAPPING_NO_VERSION] ]) recommend_opset_version = -1 for op_type in unsupported_op_types: opsets = OpMapper.OPSETS[op_type] if min(opsets.keys()) > recommend_opset_version: recommend_opset_version = min(opsets.keys()) warning_info = "\nThere are {} ops that are not supported in opset version {}, please set opset version >= {}.\n".format( len(unsupported_op_types), opset_version, recommend_opset_version) for op_type in unsupported_op_types: warning_info += "=========== {} ===========\n".format(op_type) if for_check: logging.warning(warning_info) return recommend_opset_version raise NotImplementedError(warning_info) return opset_version class CustomPaddleOp(object): CREATE_TIMES = {} def __init__(self, node): self.main_program = paddle.static.Program() self.startup_program = paddle.static.Program() self.inputs = self.create_place_holder(node) self.node = node def generate_scope_name(self, node): if node.type in CustomPaddleOp.CREATE_TIMES: CustomPaddleOp.CREATE_TIMES[node.type] += 1 else: CustomPaddleOp.CREATE_TIMES[node.type] = 1 scope_prefix = node.type + str(CustomPaddleOp.CREATE_TIMES[node.type] - 1) + '_' return scope_prefix def create_place_holder(self, node): place_holders = {} with paddle.static.program_guard(self.main_program, self.startup_program): for arg_name, idxs in node.inputs.items(): place_holders[arg_name] = [] for idx in range(len(idxs)): shape = node.input_shape(arg_name, idx) dtype = node.input_dtype(arg_name, idx) name = node.input(arg_name, idx) data = paddle.static.data( name=name, shape=shape, dtype=dtype) place_holders[arg_name].append(data) return place_holders def input(self, name, idx=None): if name not in self.inputs: return None if idx is None: return self.inputs[name] if len(self.inputs[name]) <= idx: return None return self.inputs[name][idx] def get_paddle_graph(self): scope_prefix = self.generate_scope_name(self.node) scope = paddle.static.Scope() with paddle.static.scope_guard(scope): with paddle.static.program_guard(self.main_program, self.startup_program): with paddle.utils.unique_name.guard(scope_prefix): res = self.forward() feed_var_names = [ var.name for vars in self.inputs.values() for var in vars ] fetch_vars = [var for vars in res.values() for var in vars] inference_program = graph_helper.get_program( self.main_program, feed_var_names, fetch_vars) paddle_graph = PaddleGraph.build_from_program( inference_program, feed_var_names, fetch_vars, scope=scope) output_results = dict() for arg_name, outs in res.items(): output_results[arg_name] = [out.name for out in outs] return paddle_graph, output_results def register_custom_paddle_op(paddle_op, custom_op): paddle_op_list = [] if isinstance(paddle_op, six.string_types): paddle_op_list.append(paddle_op) elif isinstance(paddle_op, list): paddle_op_list = paddle_op else: raise ValueError("paddle_op' must be List or string, but got type {}.". format(type(paddle_op))) if not isinstance(custom_op, six.class_types): raise ValueError("'custom_op' must be Class, but got type {}.".format( type(custom_op))) forward = getattr(custom_op, "forward", None) if not callable(forward): raise Exception( "Custom paddle operators must be implemented in function named 'forward'." ) for op in paddle_op_list: if op not in OpMapper.REGISTER_CUSTOM_PADDLE_OP: OpMapper.REGISTER_CUSTOM_PADDLE_OP[op] = custom_op