diff --git a/mace/python/tools/tf_converter.py b/mace/python/tools/tf_converter.py index d30a463ca2bce938d716e799f82049308e044586..8f1cbf66dfa317118227c68c61939124bf005de5 100644 --- a/mace/python/tools/tf_converter.py +++ b/mace/python/tools/tf_converter.py @@ -24,7 +24,7 @@ def main(unused_args): input_graph_def, FLAGS.input_node, FLAGS.output_node, FLAGS.prequantize) else: output_graph_def = tf_converter_lib.convert_to_mace_pb( - input_graph_def) + input_graph_def, FLAGS.runtime) with gfile.GFile(FLAGS.output, "wb") as f: f.write(output_graph_def.SerializeToString()) diff --git a/mace/python/tools/tf_converter_lib.py b/mace/python/tools/tf_converter_lib.py index 27df84accf8859a20454f4c512ce688ccea8081a..39ec8a6eb189f95a200d65a5af6383c16f4ee994 100644 --- a/mace/python/tools/tf_converter_lib.py +++ b/mace/python/tools/tf_converter_lib.py @@ -13,6 +13,12 @@ pooling_type_mode = { 'MaxPool': 2 } +buffer_type_map = { + 'FILTER' : 0, + 'IN_OUT' : 1, + 'ARGUMENT' : 2, +} + def convert_tensor(op, tensor): tf_tensor = op.outputs[0].eval() tensor.name = op.outputs[0].name @@ -30,14 +36,27 @@ def convert_tensor(op, tensor): else: raise Exception("Not supported tensor type: " + tf_dt.name) - def get_input_tensor(op, index): input_tensor = op.inputs[index] if input_tensor.op.type == 'Reshape': input_tensor = get_input_tensor(input_tensor.op, 0) return input_tensor -def convert_ops(unresolved_ops, net_def): +def add_buffer_to_image(input_name, input_type, net_def): + output_name = input_name[:-2] + "_b2i" + input_name[-2:] + op_def = net_def.op.add() + op_def.name = output_name + op_def.type = 'BufferToImage' + op_def.input.extend([input_name]) + epsilon_arg = op_def.arg.add() + epsilon_arg.name = 'buffer_type' + epsilon_arg.i = buffer_type_map[input_type] + epsilon_arg = op_def.arg.add() + epsilon_arg.name = 'mode' + epsilon_arg.i = 0 + return output_name + +def convert_ops(unresolved_ops, net_def, device): ops_count = len(unresolved_ops) resolved_count = 1 @@ -55,14 +74,13 @@ def convert_ops(unresolved_ops, net_def): op_def.type = 'DepthwiseConv2d' else: op_def.type = first_op.type - op_def.input.extend([input.name for input in first_op.inputs]) - op_def.output.extend([output.name for output in first_op.outputs]) - output_shapes = [] - for output in first_op.outputs: - output_shape = mace_pb2.OutputShape() - output_shape.dims.extend(output.shape.as_list()) - output_shapes.append(output_shape) - op_def.output_shape.extend(output_shapes) + if device == 'gpu': + op_def.input.extend([first_op.inputs[0].name]) + output_name = add_buffer_to_image(first_op.inputs[1].name, "FILTER", net_def) + op_def.input.extend([output_name]) + else: + op_def.input.extend([input.name for input in first_op.inputs]) + padding_arg = op_def.arg.add() padding_arg.name = 'padding' padding_arg.i = padding_mode[first_op.get_attr('padding')] @@ -72,11 +90,59 @@ def convert_ops(unresolved_ops, net_def): data_format_arg = op_def.arg.add() data_format_arg.name = 'data_format' data_format_arg.s = 'NHWC' + final_op = first_op + + if ops_count >= 3 and unresolved_ops[1].type == 'Const' and unresolved_ops[2].type == 'BiasAdd' : + bias_tensor = unresolved_ops[1] + tensor = net_def.tensors.add() + convert_tensor(bias_tensor, tensor) + + bias_add_op = unresolved_ops[2] + if device == 'gpu': + output_name = add_buffer_to_image(bias_add_op.inputs[1].name, "ARGUMENT", net_def) + op_def.input.extend([output_name]) + else: + op_def.input.extend([bias_add_op.inputs[1].name]) + final_op = bias_add_op + resolved_count = 3 + + if ops_count >= 4 and unresolved_ops[3].type == 'Relu': + relu_op = unresolved_ops[3]; + op_def.type = "FusedConv2D" + final_op = relu_op + resolved_count = 4 + + op_def.output.extend([output.name for output in final_op.outputs]) + output_shapes = [] + for output in final_op.outputs: + output_shape = mace_pb2.OutputShape() + output_shape.dims.extend(output.shape.as_list()) + output_shapes.append(output_shape) + op_def.output_shape.extend(output_shapes) - if ops_count >= 2 and unresolved_ops[1].type == 'BiasAdd': - bias_add_op = unresolved_ops[1] - op_def.input.extend([bias_add_op.inputs[1].name]) - resolved_count = 2 + elif first_op.type == 'FusedBatchNorm': + op_def = net_def.op.add() + op_def.name = first_op.name + op_def.type = 'BatchNorm' + if device == 'gpu': + op_def.input.extend([first_op.inputs[0].name]) + for i in range(1, len(first_op.inputs)): + output_name = add_buffer_to_image(first_op.inputs[i].name, "ARGUMENT", net_def) + op_def.input.extend([output_name]) + else: + op_def.input.extend([input.name for input in first_op.inputs]) + op_def.output.extend([first_op.outputs[0].name]) + + output_shape = mace_pb2.OutputShape() + output_shape.dims.extend(first_op.outputs[0].shape.as_list()) + op_def.output_shape.extend([output_shape]) + + epsilon_arg = op_def.arg.add() + epsilon_arg.name = 'epsilon' + epsilon_arg.f = first_op.get_attr('epsilon') + data_format_arg = op_def.arg.add() + data_format_arg.name = 'data_format' + data_format_arg.s = 'NHWC' elif first_op.type == 'Add' and first_op.name.endswith( 'batchnorm/add') and ops_count > 7: add_op = first_op @@ -90,6 +156,7 @@ def convert_ops(unresolved_ops, net_def): mul_1_op.type != 'Mul' or sub_op.type != 'Sub' or add_1_op.type != 'Add': raise Exception('Invalid BatchNorm Op') + get_input_tensor(mul_1_op, 0) input_name = get_input_tensor(mul_1_op, 0).name gamma = get_input_tensor(mul_op, 1).name beta = get_input_tensor(sub_op, 0).name @@ -168,16 +235,36 @@ def convert_ops(unresolved_ops, net_def): op_def = net_def.op.add() op_def.name = first_op.name op_def.type = "Concat" - op_def.input.extend([input.name for input in first_op.inputs]) + op_def.input.extend([first_op.inputs[i].name for i in xrange(2)]) + op_def.output.extend([output.name for output in first_op.outputs]) + axis_arg = op_def.arg.add() + axis_arg.name = 'axis' + axis_arg.i = get_input_tensor(first_op, 2).eval().astype(np.int32) + output_shapes = [] + for output in first_op.outputs: + output_shape = mace_pb2.OutputShape() + output_shape.dims.extend(output.shape.as_list()) + output_shapes.append(output_shape) + op_def.output_shape.extend(output_shapes) + elif first_op.type == 'ResizeBilinear': + op_def = net_def.op.add() + op_def.name = first_op.name + op_def.type = "ResizeBilinear" + op_def.input.extend([first_op.inputs[0].name]) op_def.output.extend([output.name for output in first_op.outputs]) + size_arg = op_def.arg.add() + size_arg.name = 'size' + size_arg.ints.extend(get_input_tensor(first_op, 1).eval().astype(np.int32).flat) + size_arg = op_def.arg.add() + size_arg.name = 'align_corners' + size_arg.ints.extend(first_op.get_attr('align_corners')) output_shapes = [] for output in first_op.outputs: output_shape = mace_pb2.OutputShape() output_shape.dims.extend(output.shape.as_list()) output_shapes.append(output_shape) op_def.output_shape.extend(output_shapes) - elif first_op.type in ['Relu', 'ResizeBilinear', 'SpaceToBatchND', - 'BatchToSpaceND', 'BiasAdd', 'FusedBatchNorm']: + elif first_op.type in ['Relu', 'SpaceToBatchND', 'BatchToSpaceND', 'BiasAdd']: op_def = net_def.op.add() op_def.name = first_op.name op_def.type = first_op.type @@ -197,7 +284,7 @@ def convert_ops(unresolved_ops, net_def): del unresolved_ops[0] -def convert_to_mace_pb(input_graph_def): +def convert_to_mace_pb(input_graph_def, device): net_def = mace_pb2.NetDef() with tf.Session() as session: @@ -206,8 +293,8 @@ def convert_to_mace_pb(input_graph_def): ops = graph.get_operations() unresolved_ops = ops while len(unresolved_ops) > 0: - convert_ops(unresolved_ops, net_def) + convert_ops(unresolved_ops, net_def, device) - print "Done." + print "PB Parsed." return net_def