未验证 提交 d1cda02c 编写于 作者: J Jason 提交者: GitHub

Merge pull request #205 from Channingss/develop

add some ops
...@@ -49,4 +49,5 @@ ...@@ -49,4 +49,5 @@
| 37 | Ceil | 38 | Concat | 39 | Flatten | 40 | ConvTranspose | | 37 | Ceil | 38 | Concat | 39 | Flatten | 40 | ConvTranspose |
| 41 | MatMul | 42 | Sum | 43 | Transpose | 44 | BatchNormalization | | 41 | MatMul | 42 | Sum | 43 | Transpose | 44 | BatchNormalization |
| 45 | Squeeze | 46 | Equal | 47 | Identity | 48 | GlobalAveragePool | | 45 | Squeeze | 46 | Equal | 47 | Identity | 48 | GlobalAveragePool |
| 49 | MaxPool | 50 | Conv | 51 | Gemm | | 49 | MaxPool | 50 | Conv | 51 | Gemm | 52 | NonZero |
| 53 | Abs | 54 | Floor |
...@@ -77,6 +77,8 @@ default_op_mapping = { ...@@ -77,6 +77,8 @@ default_op_mapping = {
'Softmax': ['softmax', ['X'], ['Out'], 'Softmax': ['softmax', ['X'], ['Out'],
dict(), dict(axis=1)], dict(), dict(axis=1)],
'Sqrt': ['sqrt', ['X'], ['Out']], 'Sqrt': ['sqrt', ['X'], ['Out']],
'Floor': ['floor', ['X'], ['Out']],
'Abs': ['abs', ['X'], ['Out']],
} }
default_ioa_constraint = { default_ioa_constraint = {
......
...@@ -31,12 +31,13 @@ from collections import OrderedDict as _dict ...@@ -31,12 +31,13 @@ from collections import OrderedDict as _dict
import math import math
import os import os
import shutil import shutil
from functools import reduce
_logger = _logging.getLogger(__name__) _logger = _logging.getLogger(__name__)
def _const_weight_or_none(node): def _const_weight_or_none(node):
if 'Constant' in node.layer_name: if 'Constant' in node.layer_type:
return node.value return node.value
if isinstance(node, ONNXGraphDataNode): if isinstance(node, ONNXGraphDataNode):
return node.weight return node.weight
...@@ -121,8 +122,13 @@ class ONNXOpMapper(OpMapper): ...@@ -121,8 +122,13 @@ class ONNXOpMapper(OpMapper):
for data_node in data_nodes: for data_node in data_nodes:
value_info = value_infos[data_node] value_info = value_infos[data_node]
ipt = np.random.random(value_info['shape']).astype( shape = value_info['shape']
value_info['dtype']) for i, dim_shape in enumerate(shape):
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'])
np.save(os.path.join(self.tmp_data_dir, data_node), ipt) np.save(os.path.join(self.tmp_data_dir, data_node), ipt)
model = onnx.shape_inference.infer_shapes(model) model = onnx.shape_inference.infer_shapes(model)
...@@ -306,9 +312,16 @@ class ONNXOpMapper(OpMapper): ...@@ -306,9 +312,16 @@ class ONNXOpMapper(OpMapper):
def place_holder(self, node): def place_holder(self, node):
self.input_shapes.append(node.out_shapes[0]) self.input_shapes.append(node.out_shapes[0])
shape = node.out_shapes[0]
for i, dim_shape in enumerate(shape):
if dim_shape == 0 and i == 0:
shape[i] = 1
if dim_shape == 0 and i != 0:
assert 'shape of input is not assigned'
attr = { attr = {
"dtype": string(node.dtype), "dtype": string(node.dtype),
"shape": node.out_shapes[0], "shape": shape,
"name": string(node.layer_name), "name": string(node.layer_name),
"append_batch_size": 'False' "append_batch_size": 'False'
} }
...@@ -406,6 +419,47 @@ class ONNXOpMapper(OpMapper): ...@@ -406,6 +419,47 @@ class ONNXOpMapper(OpMapper):
output=node, output=node,
param_attr=attr) param_attr=attr)
def RoiAlign(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True)
val_rois = self.graph.get_input_node(node, idx=1, copy=True)
pooled_height = node.get_attr('output_height')
pooled_width = node.get_attr('output_width')
spatial_scale = node.get_attr('spatial_scale')
sampling_ratio = node.get_attr('sampling_ratio')
attr = {
'pooled_height': pooled_height,
'pooled_width': pooled_width,
'spatial_scale': spatial_scale,
'sampling_ratio': sampling_ratio,
}
node.fluid_code.add_layer('roi_align',
inputs={
'input': val_x,
'rois': val_rois
},
output=node,
param_attr=attr)
def MaxRoiPool(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True)
val_rois = self.graph.get_input_node(node, idx=1, copy=True)
spatial_scale = node.get_attr('spatial_scale')
pooled_height, pooled_width = node.get_attr('pooled_shape')
attr = {
'pooled_height': pooled_height,
'pooled_width': pooled_width,
'spatial_scale': spatial_scale,
}
node.fluid_code.add_layer('roi_pool',
inputs={
'input': val_x,
'rois': val_rois
},
output=node,
param_attr=attr)
def Pad(self, node, op_independent=True): def Pad(self, node, op_independent=True):
val_x = self.graph.get_input_node(node, idx=0, copy=True) val_x = self.graph.get_input_node(node, idx=0, copy=True)
pads = node.get_attr('pads') pads = node.get_attr('pads')
...@@ -633,22 +687,23 @@ class ONNXOpMapper(OpMapper): ...@@ -633,22 +687,23 @@ class ONNXOpMapper(OpMapper):
def Slice(self, node): def Slice(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True) val_x = self.graph.get_input_node(node, idx=0, copy=True)
val_starts, val_ends, val_axes, val_steps = None, None, None, None starts, ends, axes, steps = None, None, None, None
if len(node.inputs) > 1: if len(node.inputs) > 1:
starts = self.graph.get_input_node(node, idx=1, copy=True) starts = self.graph.get_input_node(node, idx=1, copy=True)
ends = self.graph.get_input_node(node, idx=2, copy=True) ends = self.graph.get_input_node(node, idx=2, copy=True)
if len(node.inputs) > 3:
axes = self.graph.get_input_node(node, idx=3, copy=True) 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:
steps = self.graph.get_input_node(node, idx=4, copy=True) 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(starts.layer_name)
self.omit_nodes.append(ends.layer_name) self.omit_nodes.append(ends.layer_name)
self.omit_nodes.append(axes.layer_name) starts = _const_weight_or_none(starts)
self.omit_nodes.append(steps.layer_name) ends = _const_weight_or_none(ends)
starts = _const_weight_or_none(starts).copy()
ends = _const_weight_or_none(ends).copy()
axes = _const_weight_or_none(axes)
steps = _const_weight_or_none(steps)
else: else:
starts = node.get_attr('starts') starts = node.get_attr('starts')
ends = node.get_attr('ends') ends = node.get_attr('ends')
...@@ -738,6 +793,7 @@ class ONNXOpMapper(OpMapper): ...@@ -738,6 +793,7 @@ class ONNXOpMapper(OpMapper):
attr['actual_shape'] = val_shape_cast attr['actual_shape'] = val_shape_cast
else: else:
attr['actual_shape'] = val_shape attr['actual_shape'] = val_shape
if shape is None: if shape is None:
shape = val_reshaped.out_shapes[0] shape = val_reshaped.out_shapes[0]
...@@ -986,10 +1042,116 @@ class ONNXOpMapper(OpMapper): ...@@ -986,10 +1042,116 @@ class ONNXOpMapper(OpMapper):
output=node, output=node,
param_attr=attr) 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
},
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'
node.fluid_code.add_layer("cast",
inputs=not_condition,
output=cast_not_condition,
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'
node.fluid_code.add_layer("elementwise_mul",
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
},
output=mul_val_y,
param_attr=None)
node.fluid_code.add_layer("elementwise_add",
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): def Identity(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True) val_x = self.graph.get_input_node(node, idx=0, copy=True)
node.fluid_code.add_layer("assign", inputs=val_x, output=node) 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,
"name": string(node.layer_name),
}
node.fluid_code.add_layer("expand",
inputs=val_x,
output=node,
param_attr=attr)
def MaxPool(self, node): def MaxPool(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True) val_x = self.graph.get_input_node(node, idx=0, copy=True)
...@@ -1030,7 +1192,7 @@ class ONNXOpMapper(OpMapper): ...@@ -1030,7 +1192,7 @@ class ONNXOpMapper(OpMapper):
output=node, output=node,
param_attr=attr) param_attr=attr)
def GlobalAveragePool(self, node): def _global_pool(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True) val_x = self.graph.get_input_node(node, idx=0, copy=True)
val_y = self.graph.get_node(node.layer.output[0], copy=True) val_y = self.graph.get_node(node.layer.output[0], copy=True)
input_shape = val_x.out_shapes[0] input_shape = val_x.out_shapes[0]
...@@ -1042,8 +1204,15 @@ class ONNXOpMapper(OpMapper): ...@@ -1042,8 +1204,15 @@ class ONNXOpMapper(OpMapper):
poolnd = len(output_shape) - 2 # NC... poolnd = len(output_shape) - 2 # NC...
assert 2 <= poolnd <= 3, 'only pool2d and pool3d is supported' assert 2 <= poolnd <= 3, 'only pool2d and pool3d is supported'
fluid_op = 'pool{}d'.format(poolnd) fluid_op = 'pool{}d'.format(poolnd)
pool_type = None
if node.layer.op_type == 'GlobalMaxPool':
pool_type = 'max'
elif node.layer.op_type == 'GlobalAveragePool':
pool_type = 'avg'
attr = { attr = {
"pool_type": string("avg"), "pool_type": string(pool_type),
"global_pooling": True, "global_pooling": True,
"name": string(node.layer_name) "name": string(node.layer_name)
} }
...@@ -1052,6 +1221,12 @@ class ONNXOpMapper(OpMapper): ...@@ -1052,6 +1221,12 @@ class ONNXOpMapper(OpMapper):
output=node, output=node,
param_attr=attr) param_attr=attr)
def GlobalMaxPool(self, node):
self._global_pool(node)
def GlobalAveragePool(self, node):
self._global_pool(node)
def Conv(self, node): def Conv(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True) 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_w = self.graph.get_input_node(node, idx=1, copy=True)
...@@ -1108,10 +1283,11 @@ class ONNXOpMapper(OpMapper): ...@@ -1108,10 +1283,11 @@ class ONNXOpMapper(OpMapper):
def ConvTranspose(self, node): def ConvTranspose(self, node):
val_x = self.graph.get_input_node(node, idx=0, copy=True) 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_w = self.graph.get_input_node(node, idx=1, copy=True)
val_b = None
if len(node.layer.input) > 2:
val_b = self.graph.get_input_node(node, idx=2, copy=True) val_b = self.graph.get_input_node(node, idx=2, copy=True)
self.omit_nodes.append(val_w.layer_name)
self.omit_nodes.append(val_b.layer_name) self.omit_nodes.append(val_b.layer_name)
self.omit_nodes.append(val_w.layer_name)
val_y = self.graph.get_node(node.layer.output[0], copy=True) val_y = self.graph.get_node(node.layer.output[0], copy=True)
...@@ -1149,10 +1325,200 @@ class ONNXOpMapper(OpMapper): ...@@ -1149,10 +1325,200 @@ class ONNXOpMapper(OpMapper):
'dilation': dilations, 'dilation': dilations,
'groups': num_groups, 'groups': num_groups,
'param_attr': string(val_w.layer_name), 'param_attr': string(val_w.layer_name),
'bias_attr': string(val_b.layer_name), 'bias_attr': None if val_b is None else string(val_b.layer_name),
'name': string(node.layer_name), 'name': string(node.layer_name),
} }
node.fluid_code.add_layer(fluid_op, node.fluid_code.add_layer(fluid_op,
inputs=val_x, inputs=val_x,
output=node, output=node,
param_attr=attr) param_attr=attr)
def GRU(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)
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] != '':
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)
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)
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'
hidden_size = node.get_attr('hidden_size', None)
if hidden_size is None:
r_shape = val_r.out_shapes[0]
if r_shape:
hidden_size = r_shape[-1]
if hidden_size is None:
w_shape = var_w.out_shapes[0]
if w_shape:
hidden_size = w_shape[-2] // 3
if hidden_size is None and val_b:
b_shape = val_b.out_shapes[0]
if b_shape:
hidden_size = b_shape[-1] // 6
if hidden_size is None and val_xh:
xh_shape = val_xh.out_shapes[0]
if xh_shape:
hidden_size = xh_shape[-1]
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'
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)
})
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)
})
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
},
output=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'
node.fluid_code.add_layer('transpose',
inputs=var_r0,
output=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')
})
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)
})
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')
})
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)
})
var_y00 = node.layer_name + '_y00'
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,
}
node.fluid_code.add_layer('dynamic_gru',
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] != '':
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] != '':
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])
})
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册