From f05ae7377bef2919ea2986bb2862888bd8f0edc7 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 22 Nov 2019 02:39:32 +0000 Subject: [PATCH] add support for OP:NonZero --- x2paddle/op_mapper/onnx_op_mapper.py | 441 ++++++++++++++++----------- 1 file changed, 268 insertions(+), 173 deletions(-) diff --git a/x2paddle/op_mapper/onnx_op_mapper.py b/x2paddle/op_mapper/onnx_op_mapper.py index 90b50ff..8ea62ca 100644 --- a/x2paddle/op_mapper/onnx_op_mapper.py +++ b/x2paddle/op_mapper/onnx_op_mapper.py @@ -31,7 +31,7 @@ from collections import OrderedDict as _dict import math import os import shutil - +from functools import reduce _logger = _logging.getLogger(__name__) @@ -51,13 +51,14 @@ def get_same_padding(in_size, kernel_size, stride): return [pad0, pad1] -class ONNXOpMapper(OpMapper): +class ONNXOpMapper(OpMapper): elementwise_ops = { 'Add': 'elementwise_add', 'Div': 'elementwise_div', 'Sub': 'elementwise_sub', 'Mul': 'elementwise_mul', - 'Pow': 'elementwise_pow',} + 'Pow': 'elementwise_pow', + } def __init__(self, decoder, save_dir): super(ONNXOpMapper, self).__init__() @@ -70,10 +71,10 @@ class ONNXOpMapper(OpMapper): self.is_inference = False self.tmp_data_dir = os.path.join(save_dir, 'tmp_data') self.get_output_shapes() - + if not self.op_checker(): raise Exception("Model are not supported yet.") - + #mapping op print("Total nodes: {}".format( sum([ @@ -117,19 +118,18 @@ class ONNXOpMapper(OpMapper): def get_results_of_inference(self, model, value_infos, data_nodes): if not os.path.exists(self.tmp_data_dir): os.makedirs(self.tmp_data_dir) - + for data_node in data_nodes: value_info = value_infos[data_node] shape = value_info['shape'] for i, dim_shape in enumerate(shape): - if dim_shape==0 and i==0: - shape[i]=1 - if dim_shape==0 and i!=0: + if dim_shape == 0 and i == 0: + shape[i] = 1 + if dim_shape == 0 and i != 0: assert 'shape of input is not assigned' - ipt = np.random.random(shape).astype( - value_info['dtype']) + ipt = np.random.random(shape).astype(value_info['dtype']) np.save(os.path.join(self.tmp_data_dir, data_node), ipt) - + model = onnx.shape_inference.infer_shapes(model) outputs = [] for value_info in model.graph.value_info: @@ -165,8 +165,8 @@ class ONNXOpMapper(OpMapper): for opt in layer.output: if opt in value_infos: value_info = value_infos[opt] - if len(value_info['shape'] - ) == 0 or value_info['dtype'] is None or 0 in value_info['shape']: + if len(value_info['shape']) == 0 or value_info[ + 'dtype'] is None or 0 in value_info['shape']: if self.is_inference == False: self.get_results_of_inference( onnx_model, value_infos, @@ -263,16 +263,17 @@ class ONNXOpMapper(OpMapper): if child_func_code is not None: self.used_custom_layers[op + '_child_func'] = child_func_code + def elementwise_map(self, node): assert node.layer_type in self.elementwise_ops op_type = self.elementwise_ops[node.layer_type] - + val_x = self.graph.get_input_node(node, idx=0, copy=True) val_y = self.graph.get_input_node(node, idx=1, copy=True) val_y_shape = val_y.out_shapes[0] val_x_shape = val_x.out_shapes[0] - - if len(val_x_shape)2 " - if axis==0 and len(indices_shape)<=1: + if axis == 0 and len(indices_shape) <= 1: node.fluid_code.add_layer('gather', inputs={ 'input': val_x, @@ -639,14 +644,16 @@ class ONNXOpMapper(OpMapper): inputs=node, output=node, param_attr=attr_trans) - elif len(indices_shape)>1: + elif len(indices_shape) > 1: from functools import reduce - reshape_shape = reduce(lambda x,y:x*y, indices_shape) + reshape_shape = reduce(lambda x, y: x * y, indices_shape) node.fluid_code.add_layer('reshape', inputs=indices, output=indices, - param_attr={'shape':[reshape_shape,]}) - + param_attr={'shape': [ + reshape_shape, + ]}) + perm = list(range(len(val_x.out_shapes[0]))) perm = [axis] + perm[:axis] + perm[axis + 1:] attr_trans = {'perm': perm} @@ -675,7 +682,7 @@ class ONNXOpMapper(OpMapper): node.fluid_code.add_layer('reshape', inputs=node, output=node, - param_attr={'shape':reshaped_shape}) + param_attr={'shape': reshaped_shape}) def Slice(self, node): val_x = self.graph.get_input_node(node, idx=0, copy=True) @@ -683,15 +690,15 @@ class ONNXOpMapper(OpMapper): if len(node.inputs) > 1: starts = self.graph.get_input_node(node, idx=1, copy=True) ends = self.graph.get_input_node(node, idx=2, copy=True) - if len(node.inputs)>3: + if len(node.inputs) > 3: axes = self.graph.get_input_node(node, idx=3, copy=True) self.omit_nodes.append(axes.layer_name) axes = _const_weight_or_none(axes) - if len(node.inputs)>4: + if len(node.inputs) > 4: steps = self.graph.get_input_node(node, idx=4, copy=True) self.omit_nodes.append(steps.layer_name) steps = _const_weight_or_none(steps) - + self.omit_nodes.append(starts.layer_name) self.omit_nodes.append(ends.layer_name) starts = _const_weight_or_none(starts) @@ -756,7 +763,7 @@ class ONNXOpMapper(OpMapper): 'dim': axis, 'name': string(node.layer_name) } - + node.fluid_code.add_layer('split', inputs=val_x, output=val_y, @@ -770,18 +777,18 @@ class ONNXOpMapper(OpMapper): if isinstance(val_shape, ONNXGraphDataNode): self.omit_nodes.append(val_shape.layer_name) - + attr = {'name': string(node.layer_name)} # catch dynamic graph shape if isinstance(val_shape, ONNXGraphNode): shape, _, _ = self.get_dynamic_shape(val_shape.layer_name) if val_shape.dtype == 'int64': - val_shape_cast = val_shape.layer_name+'_cast' + val_shape_cast = val_shape.layer_name + '_cast' node.fluid_code.add_layer('cast', - inputs=val_shape, - output=val_shape_cast, - param_attr={'dtype':string('int32')}) - + inputs=val_shape, + output=val_shape_cast, + param_attr={'dtype': string('int32')}) + attr['actual_shape'] = val_shape_cast else: attr['actual_shape'] = val_shape @@ -796,7 +803,7 @@ class ONNXOpMapper(OpMapper): 'input "shape" not inferred, use [1, -1] as dummy value, ' 'the behavior of Paddle fluid maybe undefined', node.layer_name, val_x.layer_name, val_reshaped.layer_name) - + attr['shape'] = shape node.fluid_code.add_layer('reshape', inputs=val_x, @@ -1033,73 +1040,117 @@ class ONNXOpMapper(OpMapper): inputs=val_x, output=node, param_attr=attr) - + def Equal(self, node): val_x = self.graph.get_input_node(node, idx=0, copy=True) val_y = self.graph.get_input_node(node, idx=1, copy=True) node.fluid_code.add_layer("equal", - inputs={'x':val_x, 'y':val_y}, + inputs={ + 'x': val_x, + 'y': val_y + }, output=node, param_attr=None) + def Where(self, node): condition = self.graph.get_input_node(node, idx=0, copy=True) val_x = self.graph.get_input_node(node, idx=1, copy=True) val_y = self.graph.get_input_node(node, idx=2, copy=True) - + not_condition = condition.layer_name + '_not' node.fluid_code.add_layer("logical_not", inputs=condition, output=not_condition, param_attr=None) - cast_not_condition = not_condition+'_cast' + cast_not_condition = not_condition + '_cast' node.fluid_code.add_layer("cast", inputs=not_condition, output=cast_not_condition, - param_attr={'dtype':string(val_x.dtype)}) + param_attr={'dtype': string(val_x.dtype)}) cast_condition = condition.layer_name + '_cast' node.fluid_code.add_layer("cast", inputs=condition, output=cast_condition, - param_attr={'dtype':string(val_x.dtype)}) - mul_val_x = val_x.layer_name + '_mul' + param_attr={'dtype': string(val_x.dtype)}) + mul_val_x = val_x.layer_name + '_mul' node.fluid_code.add_layer("elementwise_mul", - inputs={'x':val_x,'y':cast_condition}, + inputs={ + 'x': val_x, + 'y': cast_condition + }, output=mul_val_x, param_attr=None) - + mul_val_y = val_y.layer_name + '_mul' node.fluid_code.add_layer("elementwise_mul", - inputs={'x':val_y,'y':cast_not_condition}, + inputs={ + 'x': val_y, + 'y': cast_not_condition + }, output=mul_val_y, param_attr=None) - + node.fluid_code.add_layer("elementwise_add", - inputs={'x':mul_val_x,'y':mul_val_y}, + inputs={ + 'x': mul_val_x, + 'y': mul_val_y + }, output=node, param_attr=None) - + + def NonZero(self, node): + val_x = self.graph.get_input_node(node, idx=0, copy=True) + where_name = node.layer_name + '_where' + node.fluid_code.add_layer("where", + inputs=val_x.layer_name + '==1', + output=where_name) + dims = len(val_x.out_shapes[0]) + elements_count_val_x = reduce(lambda x, y: x * y, val_x.out_shapes[0]) + flatten_names = [] + for dim in range(dims): + slice_name = node.layer_name + '_slice' + str(dim) + flatten_name = node.layer_name + '_flatten' + str(dim) + flatten_names.append(flatten_name) + attr = { + 'axes': list(range(dims)), + 'starts': [0, dim], + 'ends': [elements_count_val_x, dim + 1] + } + node.fluid_code.add_layer("slice", + inputs=where_name, + output=slice_name, + param_attr=attr) + node.fluid_code.add_layer("flatten", + inputs=slice_name, + output=flatten_name, + param_attr={'axis': 0}) + node.fluid_code.add_layer("concat", + inputs=flatten_names, + output=node, + param_attr={'axis': 0}) + def Identity(self, node): val_x = self.graph.get_input_node(node, idx=0, copy=True) node.fluid_code.add_layer("assign", inputs=val_x, output=node) - + def Tile(self, node): val_x = self.graph.get_input_node(node, idx=0, copy=True) val_repeats = self.graph.get_input_node(node, idx=1, copy=True) repeats = _const_weight_or_none(val_repeats) assert repeats is not None, 'for OP:Tile, only const repeats supported' - + if isinstance(repeats, int): repeats = [repeats] - + attr = { - 'expand_times':repeats, + 'expand_times': repeats, "name": string(node.layer_name), } - node.fluid_code.add_layer("expand", - inputs=val_x, + node.fluid_code.add_layer("expand", + inputs=val_x, output=node, param_attr=attr) - + def MaxPool(self, node): val_x = self.graph.get_input_node(node, idx=0, copy=True) @@ -1152,7 +1203,7 @@ class ONNXOpMapper(OpMapper): poolnd = len(output_shape) - 2 # NC... assert 2 <= poolnd <= 3, 'only pool2d and pool3d is supported' fluid_op = 'pool{}d'.format(poolnd) - + pool_type = None if node.layer.op_type == 'GlobalMaxPool': pool_type = 'max' @@ -1168,13 +1219,13 @@ class ONNXOpMapper(OpMapper): inputs=val_x, output=node, param_attr=attr) - + def GlobalMaxPool(self, node): self._global_pool(node) - + def GlobalAveragePool(self, node): self._global_pool(node) - + def Conv(self, node): val_x = self.graph.get_input_node(node, idx=0, copy=True) val_w = self.graph.get_input_node(node, idx=1, copy=True) @@ -1232,7 +1283,7 @@ class ONNXOpMapper(OpMapper): val_x = self.graph.get_input_node(node, idx=0, copy=True) val_w = self.graph.get_input_node(node, idx=1, copy=True) val_b = None - if len(node.layer.input)>2: + if len(node.layer.input) > 2: val_b = self.graph.get_input_node(node, idx=2, copy=True) self.omit_nodes.append(val_b.layer_name) self.omit_nodes.append(val_w.layer_name) @@ -1285,27 +1336,31 @@ class ONNXOpMapper(OpMapper): val_x = self.graph.get_input_node(node, idx=0, copy=True) val_w = self.graph.get_input_node(node, idx=1, copy=True) val_r = self.graph.get_input_node(node, idx=2, copy=True) - + val_b = None val_len = None val_xh = None miss_arg_num = 0 num_ipt = len(node.layer.input) - if num_ipt>3 and node.layer.input[3] != '': + if num_ipt > 3 and node.layer.input[3] != '': val_b = self.graph.get_input_node(node, idx=3, copy=True) else: miss_arg_num += 1 - if num_ipt>4 and node.layer.input[4] != '': - val_len = self.graph.get_input_node(node, idx=4-miss_arg_num, copy=True) + if num_ipt > 4 and node.layer.input[4] != '': + val_len = self.graph.get_input_node(node, + idx=4 - miss_arg_num, + copy=True) else: miss_arg_num += 1 - if num_ipt>5 and node.layer.input[5] != '': - val_xh = self.graph.get_input_node(node, idx=5-miss_arg_num, copy=True) - + if num_ipt > 5 and node.layer.input[5] != '': + val_xh = self.graph.get_input_node(node, + idx=5 - miss_arg_num, + copy=True) + data, dtype, shape = self.get_dynamic_shape(val_x.layer_name) - + x_shape = val_x.out_shapes[0] - + assert x_shape[1] == 1, 'only X with batch_size = 1 supported' assert node.get_attr('clip', None) is None, 'clipping not supported' @@ -1326,103 +1381,143 @@ class ONNXOpMapper(OpMapper): xh_shape = val_xh.out_shapes[0] if xh_shape: hidden_size = xh_shape[-1] - - direction = node.get_attr('direction', 'forward') + + direction = node.get_attr('direction', 'forward') assert direction != 'bidirectional', 'direction = bidirectional not supported' - + activations = node.get_attr('activations', ['Sigmoid', 'Tanh']) assert len(activations) == 2, 'bidirectional operation not supported' - - assert node.get_attr( - 'linear_before_reset', - 0) == 0, 'only linear_before_reset = 0 supported' - + + assert node.get_attr('linear_before_reset', + 0) == 0, 'only linear_before_reset = 0 supported' + activations = [s.lower() for s in activations] gate_activation, candidate_activation = activations is_reverse = direction == 'reverse' - + var_x0 = node.layer_name + '_x0' node.fluid_code.add_layer('squeeze', inputs=val_x, output=var_x0, - param_attr={'axes': [1],'name':string(var_x0)}) - + param_attr={ + 'axes': [1], + 'name': string(var_x0) + }) + var_w0 = node.layer_name + '_w0' node.fluid_code.add_layer('squeeze', inputs=val_w, output=var_w0, - param_attr={'axes': [0],'name':string(var_w0)}) - + param_attr={ + 'axes': [0], + 'name': string(var_w0) + }) + var_fc = node.layer_name + '_fc' var_mm = (node.layer_name + '_mm') if val_b else var_fc node.fluid_code.add_layer('matmul', - inputs={'x':var_x0, 'y':var_w0}, + inputs={ + 'x': var_x0, + 'y': var_w0 + }, output=var_mm, - param_attr={'transpose_x': 0,'transpose_y': 1,'name':string(var_mm)}) - + param_attr={ + 'transpose_x': 0, + 'transpose_y': 1, + 'name': string(var_mm) + }) + var_r0 = node.layer_name + '_r0' node.fluid_code.add_layer('squeeze', inputs=val_r, output=var_r0, - param_attr={'axes': [0],'name':string(var_r0)}) - - var_r0t = node.layer_name + '_r0t' - + param_attr={ + 'axes': [0], + 'name': string(var_r0) + }) + + var_r0t = node.layer_name + '_r0t' + node.fluid_code.add_layer('transpose', inputs=var_r0, output=var_r0t, - param_attr={'perm': [1, 0],'name':string(var_r0t)}) + param_attr={ + 'perm': [1, 0], + 'name': string(var_r0t) + }) if val_b: var_bi = node.layer_name + '_bi' var_bh = node.layer_name + '_bh' node.fluid_code.add_layer('split', - inputs=val_b, - output=var_bi+','+var_bh, - param_attr={'axis': 1, - 'split': [hidden_size * 3, hidden_size * 3], - 'name':string(node.layer_name+'.b/split')}) + inputs=val_b, + output=var_bi + ',' + var_bh, + param_attr={ + 'axis': + 1, + 'split': + [hidden_size * 3, hidden_size * 3], + 'name': + string(node.layer_name + '.b/split') + }) var_bi0 = node.layer_name + '_bi0' node.fluid_code.add_layer('squeeze', - inputs=var_bi, - output=var_bi0, - param_attr={'axes': [0],'name':string(var_bi0)}) - + inputs=var_bi, + output=var_bi0, + param_attr={ + 'axes': [0], + 'name': string(var_bi0) + }) + node.fluid_code.add_layer('elmentwise_add', - inputs=[var_mm, var_bi0], - output=var_fc, - param_attr={'axes': 1,'name':string(node.layer_name+'.i/bias')}) + inputs=[var_mm, var_bi0], + output=var_fc, + param_attr={ + 'axes': + 1, + 'name': + string(node.layer_name + '.i/bias') + }) if val_xh: var_xh0 = node.layer_name + '_xh0' node.fluid_code.add_layer('squeeze', - inputs=val_xh, - output=var_xh0, - param_attr={'axes': [1],'name':string(var_xh0)}) + inputs=val_xh, + output=var_xh0, + param_attr={ + 'axes': [1], + 'name': string(var_xh0) + }) var_y00 = node.layer_name + '_y00' - - attr={ - 'origin_mode':True, + + attr = { + 'origin_mode': True, 'h_0': var_xh0 if val_xh else None, - 'is_reverse':is_reverse, - 'gate_activation':string(gate_activation), - 'candidate_activation':string(candidate_activation), - 'param_attr':string(var_r0t), - 'bias_attr':string(var_bh) if val_b else False, + 'is_reverse': is_reverse, + 'gate_activation': string(gate_activation), + 'candidate_activation': string(candidate_activation), + 'param_attr': string(var_r0t), + 'bias_attr': string(var_bh) if val_b else False, } node.fluid_code.add_layer('dynamic_gru', - inputs=var_fc +','+ str(hidden_size), + inputs=var_fc + ',' + str(hidden_size), output=var_y00, param_attr=attr) - + num_opt = len(node.layer.output) - - if num_opt>0 and node.layer.output[0] != '': + + if num_opt > 0 and node.layer.output[0] != '': node.fluid_code.add_layer('unsqueeze', - inputs=var_y00, - output=node.layer.output[0], - param_attr={'axes': [1, 1],'name':string(node.layer.output[0])}) - if num_opt>1 and node.layer.output[1] != '': + inputs=var_y00, + output=node.layer.output[0], + param_attr={ + 'axes': [1, 1], + 'name': string(node.layer.output[0]) + }) + if num_opt > 1 and node.layer.output[1] != '': node.fluid_code.add_layer('unsqueeze', - inputs=var_y00, - output=node.layer.output[1], - param_attr={'axes': [1, 1],'name':string(node.layer.output[1])}) + inputs=var_y00, + output=node.layer.output[1], + param_attr={ + 'axes': [1, 1], + 'name': string(node.layer.output[1]) + }) -- GitLab