From 66a68689096716913b1016fec2095f4be66c85f6 Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Fri, 8 Dec 2017 15:27:14 +0800
Subject: [PATCH] Refactor tf_converter_lib for robust.

---
 mace/python/tools/tf_converter_lib.py | 652 +++++++++++++++-----------
 1 file changed, 366 insertions(+), 286 deletions(-)

diff --git a/mace/python/tools/tf_converter_lib.py b/mace/python/tools/tf_converter_lib.py
index 08d9dcd7..85e39673 100644
--- a/mace/python/tools/tf_converter_lib.py
+++ b/mace/python/tools/tf_converter_lib.py
@@ -24,22 +24,10 @@ data_type_map = {
   'DT_FLOAT': mace_pb2.DT_FLOAT
 }
 
-def convert_tensor(op, tensor):
-  tf_tensor = op.outputs[0].eval()
-  tensor.name = op.outputs[0].name
-
-  shape = list(tf_tensor.shape)
-  tensor.dims.extend(shape)
-
-  tf_dt = op.get_attr('dtype')
-  if tf_dt == tf.float32:
-    tensor.data_type = mace_pb2.DT_FLOAT
-    tensor.float_data.extend(tf_tensor.astype(float).flat)
-  elif tf_dt == tf.int32:
-    tensor.data_type = mace_pb2.DT_INT32
-    tensor.int32_data.extend(tf_tensor.astype(np.int32).flat)
-  else:
-    raise Exception("Not supported tensor type: " + tf_dt.name)
+BATCH_NORM_ORDER = ["Add", "Rsqrt", "Mul", "Mul", "Mul", "Sub", "Add"]
+
+MACE_INPUT_NODE_NAME = "mace_input_node"
+MACE_OUTPUT_NODE_NAME = "mace_output_node"
 
 def get_input_tensor(op, index):
   input_tensor = op.inputs[index]
@@ -47,281 +35,378 @@ def get_input_tensor(op, index):
     input_tensor = get_input_tensor(input_tensor.op, 0)
   return input_tensor
 
-def add_buffer_to_image(input_name, input_type, dt, net_def):
-  output_name = input_name[:-2] + "_b2i" + input_name[-2:]
-  op_def = net_def.op.add()
-  op_def.name = output_name[:-2]
-  op_def.type = 'BufferToImage'
-  op_def.input.extend([input_name])
-  op_def.output.extend([output_name])
-
-  arg = op_def.arg.add()
-  arg.name = 'buffer_type'
-  arg.i = buffer_type_map[input_type]
-  arg = op_def.arg.add()
-  arg.name = 'mode'
-  arg.i = 0
-  arg = op_def.arg.add()
-  arg.name = 'T'
-  arg.i = dt
-  return output_name
-
-def add_image_to_buffer(input_name, input_type, dt, net_def):
-  output_name = input_name[:-2] + "_i2b" + input_name[-2:]
-  op_def = net_def.op.add()
-  op_def.name = output_name[:-2]
-  op_def.type = 'ImageToBuffer'
-  op_def.input.extend([input_name])
-  op_def.output.extend([output_name])
-
-  arg = op_def.arg.add()
-  arg.name = 'buffer_type'
-  arg.i = buffer_type_map[input_type]
-  arg = op_def.arg.add()
-  arg.name = 'T'
-  arg.i = dt
-  return output_name
-
-def add_input_transform(name, dt, net_def):
-  new_input_name = "mace_input_node:0"
-  op_def = net_def.op.add()
-  op_def.name = name
-  op_def.type = 'BufferToImage'
-  op_def.input.extend([new_input_name])
-  op_def.output.extend([name+':0'])
-
-  epsilon_arg = op_def.arg.add()
-  epsilon_arg.name = 'buffer_type'
-  epsilon_arg.i = buffer_type_map['IN_OUT']
-
-  arg = op_def.arg.add()
-  arg.name = 'T'
-  arg.i = dt
-
-def add_output_transform(name, net_def):
-  output_name = "mace_output_node:0"
-  op_def = net_def.op.add()
-  op_def.name = output_name[:-2]
-  op_def.type = 'ImageToBuffer'
-  op_def.input.extend([name+':0'])
-  op_def.output.extend([output_name])
-
-  epsilon_arg = op_def.arg.add()
-  epsilon_arg.name = 'buffer_type'
-  epsilon_arg.i = buffer_type_map['IN_OUT']
-
-def add_output_shape(outputs, op):
-  output_shapes = []
-  for output in outputs:
-    if output.shape is not None and not output.shape:
-      output_shape = mace_pb2.OutputShape()
-      output_shape.dims.extend(output.shape.as_list())
-      output_shapes.append(output_shape)
-  op.output_shape.extend(output_shapes)
-
-def convert_ops(unresolved_ops, dt, net_def, device):
-  ops_count = len(unresolved_ops)
-  resolved_count = 1
-
-  first_op = unresolved_ops[0]
-
-  if first_op.type in ['Placeholder', 'Reshape', 'Identity']:
-    pass
-  elif first_op.type == 'Const':
-    tensor = net_def.tensors.add()
-    convert_tensor(first_op, tensor)
-  else:
-    op_def = net_def.op.add()
+class TFConverter(object):
+  def __init__(self, tf_ops, net_def, dt, device):
+    self.net_def = net_def
+    self.tf_ops = tf_ops
+    self.dt = dt
+    self.device = device
+    self.tf_graph = {}
+    self.resolved_ops = {}
+
+    for op in tf_ops:
+      self.resolved_ops[op.name] = 0
+      for input in op.inputs:
+        input_name = input.name[:-2]
+        if input_name not in self.tf_graph:
+          self.tf_graph[input_name] = []
+        print input_name
+        self.tf_graph[input_name].append(op)
+
+  def add_buffer_to_image(self, input_name, input_type):
+    output_name = input_name[:-2] + "_b2i" + input_name[-2:]
+    op_def = self.net_def.op.add()
+    op_def.name = output_name[:-2]
+    op_def.type = 'BufferToImage'
+    op_def.input.extend([input_name])
+    op_def.output.extend([output_name])
+
+    arg = op_def.arg.add()
+    arg.name = 'buffer_type'
+    arg.i = buffer_type_map[input_type]
+    arg = op_def.arg.add()
+    arg.name = 'mode'
+    arg.i = 0
     arg = op_def.arg.add()
     arg.name = 'T'
-    arg.i = dt
+    arg.i = self.dt
+    return output_name
 
-    if first_op.type == 'Conv2D' or first_op.type == 'DepthwiseConv2dNative':
-      op_def.name = first_op.name
-      if first_op.type == 'DepthwiseConv2dNative':
-        op_def.type = 'DepthwiseConv2d'
-      else:
-        op_def.type = first_op.type
-      if device == 'gpu':
-        op_def.input.extend([first_op.inputs[0].name])
-        output_name = add_buffer_to_image(first_op.inputs[1].name, "FILTER", dt, net_def)
+  def add_input_transform(self, name):
+    new_input_name = MACE_INPUT_NODE_NAME + ":0"
+    op_def = self.net_def.op.add()
+    op_def.name = name
+    op_def.type = 'BufferToImage'
+    op_def.input.extend([new_input_name])
+    op_def.output.extend([name+':0'])
+
+    epsilon_arg = op_def.arg.add()
+    epsilon_arg.name = 'buffer_type'
+    epsilon_arg.i = buffer_type_map['IN_OUT']
+
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+
+  def add_output_transform(self, name):
+    output_name = MACE_OUTPUT_NODE_NAME + ":0"
+    op_def = self.net_def.op.add()
+    op_def.name = output_name[:-2]
+    op_def.type = 'ImageToBuffer'
+    op_def.input.extend([name+':0'])
+    op_def.output.extend([output_name])
+
+    epsilon_arg = op_def.arg.add()
+    epsilon_arg.name = 'buffer_type'
+    epsilon_arg.i = buffer_type_map['IN_OUT']
+
+  @staticmethod
+  def add_output_shape(outputs, op):
+    output_shapes = []
+    for output in outputs:
+      if output.shape is not None and not output.shape:
+        output_shape = mace_pb2.OutputShape()
+        output_shape.dims.extend(output.shape.as_list())
+        output_shapes.append(output_shape)
+    op.output_shape.extend(output_shapes)
+
+  def convert_tensor(self, op):
+    tensor = self.net_def.tensors.add()
+    tf_tensor = op.outputs[0].eval()
+    tensor.name = op.outputs[0].name
+
+    shape = list(tf_tensor.shape)
+    tensor.dims.extend(shape)
+
+    tf_dt = op.get_attr('dtype')
+    if tf_dt == tf.float32:
+      tensor.data_type = mace_pb2.DT_FLOAT
+      tensor.float_data.extend(tf_tensor.astype(np.float32).flat)
+    elif tf_dt == tf.int32:
+      tensor.data_type = mace_pb2.DT_INT32
+      tensor.int32_data.extend(tf_tensor.astype(np.int32).flat)
+    else:
+      raise Exception("Not supported tensor type: " + tf_dt.name)
+    self.resolved_ops[op.name] = 1
+
+  def convert_conv2d(self, op):
+    op_def = mace_pb2.OperatorDef()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    if op.type == 'DepthwiseConv2dNative':
+      op_def.type = 'DepthwiseConv2d'
+    else:
+      op_def.type = op.type
+    if self.device == 'gpu':
+      op_def.input.extend([op.inputs[0].name])
+      output_name = self.add_buffer_to_image(op.inputs[1].name, "FILTER")
+      op_def.input.extend([output_name])
+    else:
+      op_def.input.extend([input.name for input in op.inputs])
+
+    padding_arg = op_def.arg.add()
+    padding_arg.name = 'padding'
+    padding_arg.i = padding_mode[op.get_attr('padding')]
+    strides_arg = op_def.arg.add()
+    strides_arg.name = 'strides'
+    strides_arg.ints.extend(op.get_attr('strides')[1:3])
+    data_format_arg = op_def.arg.add()
+    data_format_arg.name = 'data_format'
+    data_format_arg.s = 'NHWC'
+    final_op = op
+    self.resolved_ops[op.name] = 1
+
+    if len(self.tf_graph[op.name]) == 1 and self.tf_graph[op.name][0].type == 'BiasAdd' :
+      bias_add_op = self.tf_graph[op.name][0]
+      if self.device == 'gpu':
+        output_name = self.add_buffer_to_image(bias_add_op.inputs[1].name, "ARGUMENT")
         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')]
-      strides_arg = op_def.arg.add()
-      strides_arg.name = 'strides'
-      strides_arg.ints.extend(first_op.get_attr('strides')[1:3])
-      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", dt, 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])
-      add_output_shape(final_op.outputs, op_def)
-
-    elif first_op.type == 'FusedBatchNorm':
-      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", dt, net_def)
-          op_def.input.extend([output_name])
+        op_def.input.extend([bias_add_op.inputs[1].name])
+      final_op = bias_add_op
+      self.resolved_ops[bias_add_op.name] = 1
+
+    if len(self.tf_graph[final_op.name]) == 1 \
+        and self.tf_graph[final_op.name][0].type == 'Relu':
+      relu_op = self.tf_graph[final_op.name][0]
+      op_def.type = "FusedConv2D"
+      final_op = relu_op
+      self.resolved_ops[relu_op.name] = 1
+
+    op_def.output.extend([output.name for output in final_op.outputs])
+    self.add_output_shape(final_op.outputs, op_def)
+    self.net_def.op.extend([op_def])
+
+  def convert_fused_batchnorm(self, op):
+    op_def = mace_pb2.OperatorDef()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = 'BatchNorm'
+    if self.device == 'gpu':
+      op_def.input.extend([op.inputs[0].name])
+      for i in range(1, len(op.inputs)):
+        output_name = self.add_buffer_to_image(op.inputs[i].name, "ARGUMENT")
+        op_def.input.extend([output_name])
+    else:
+      op_def.input.extend([input.name for input in op.inputs])
+    op_def.output.extend([op.outputs[0].name])
+
+    self.add_output_shape(op.outputs, op_def)
+
+    epsilon_arg = op_def.arg.add()
+    epsilon_arg.name = 'epsilon'
+    epsilon_arg.f = op.get_attr('epsilon')
+    data_format_arg = op_def.arg.add()
+    data_format_arg.name = 'data_format'
+    data_format_arg.s = 'NHWC'
+    self.resolved_ops[op.name] = 1
+    self.net_def.op.extend([op_def])
+
+  def convert_batchnorm(self, op):
+    bn_ops = []
+    bn_ops.append(op)
+    for i in range(1, 7):
+      if len(self.tf_graph[bn_ops[i-1].name]) == 1 \
+          and self.tf_graph[bn_ops[i-1].name][0].type == BATCH_NORM_ORDER[i]:
+        bn_ops.append(self.tf_graph[bn_ops[i-1].name][0])
       else:
-        op_def.input.extend([input.name for input in first_op.inputs])
-      op_def.output.extend([first_op.outputs[0].name])
-
-      add_output_shape(first_op.outputs, op_def)
-
-      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
-      mul_op = unresolved_ops[2]
-      mul_1_op = unresolved_ops[3]
-      mul_2_op = unresolved_ops[4]
-      sub_op = unresolved_ops[5]
-      add_1_op = unresolved_ops[6]
-      print (mul_op.type, mul_2_op.type, mul_1_op.type, sub_op.type)
-      if mul_op.type != 'Mul' or mul_2_op.type != 'Mul' or \
-                      mul_1_op.type != 'Mul' or sub_op.type != 'Sub' or add_1_op.type != 'Add':
         raise Exception('Invalid BatchNorm Op')
 
-      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
-      mean = get_input_tensor(mul_2_op, 0).name
-      variance = get_input_tensor(add_op, 0).name
-
-      op_def.name = first_op.name[:-4]  # remove /add
-      op_def.type = 'BatchNorm'
-      if device == 'gpu':
-        op_def.input.extend([input_name])
-        for tensor_name in [gamma, beta, mean, variance]:
-          output_name = add_buffer_to_image(tensor_name, "ARGUMENT", dt, net_def)
-          op_def.input.extend([output_name])
-      else:
-        op_def.input.extend([input_name, gamma, beta, mean, variance])
-      op_def.output.extend([output.name for output in add_1_op.outputs])
-      add_output_shape(add_1_op.outputs, op_def)
-      epsilon_arg = op_def.arg.add()
-      epsilon_arg.name = 'epsilon'
-      epsilon_arg.f = get_input_tensor(add_op, 1).eval().astype(np.float)
-      data_format_arg = op_def.arg.add()
-      data_format_arg.name = 'data_format'
-      data_format_arg.s = 'NHWC'
-
-      resolved_count = 7
-    elif first_op.type == 'Relu6':
-      op_def.name = first_op.name
-      op_def.type = 'Relu'
-      op_def.input.extend([input.name for input in first_op.inputs])
-      op_def.output.extend([output.name for output in first_op.outputs])
-      add_output_shape(first_op.outputs, op_def)
-      max_limit_arg = op_def.arg.add()
-      max_limit_arg.name = 'max_limit'
-      max_limit_arg.f = 6
-    elif first_op.type == 'AvgPool' or first_op.type == 'MaxPool':
-      op_def.name = first_op.name
-      op_def.type = 'Pooling'
-      op_def.input.extend([input.name for input in first_op.inputs])
-      op_def.output.extend([output.name for output in first_op.outputs])
-      add_output_shape(first_op.outputs, op_def)
-      pooling_type_arg = op_def.arg.add()
-      pooling_type_arg.name = 'pooling_type'
-      pooling_type_arg.i = pooling_type_mode[first_op.type]
-      padding_arg = op_def.arg.add()
-      padding_arg.name = 'padding'
-      padding_arg.i = padding_mode[first_op.get_attr('padding')]
-      strides_arg = op_def.arg.add()
-      strides_arg.name = 'strides'
-      strides_arg.ints.extend(first_op.get_attr('strides')[1:3])
-      kernels_arg = op_def.arg.add()
-      kernels_arg.name = 'kernels'
-      kernels_arg.ints.extend(first_op.get_attr('ksize')[1:3])
-      data_format_arg = op_def.arg.add()
-      data_format_arg.name = 'data_format'
-      data_format_arg.s = 'NHWC'
-    elif first_op.type == 'Add':
-      op_def.name = first_op.name
-      op_def.type = "AddN"
-      op_def.input.extend([input.name for input in first_op.inputs])
-      op_def.output.extend([output.name for output in first_op.outputs])
-      add_output_shape(first_op.outputs, op_def)
-    elif first_op.type == 'ConcatV2':
-      op_def.name = first_op.name
-      op_def.type = "Concat"
-      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)
-      add_output_shape(first_op.outputs, op_def)
-    elif first_op.type == 'ResizeBilinear':
-      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.i = first_op.get_attr('align_corners')
-      add_output_shape(first_op.outputs, op_def)
-    elif first_op.type == 'BiasAdd':
-      op_def.name = first_op.name
-      op_def.type = first_op.type
-      op_def.input.extend([first_op.inputs[0].name])
-      if device == 'gpu':
-        output_name = add_buffer_to_image(first_op.inputs[1].name, "ARGUMENT", dt, net_def)
+    op_def = mace_pb2.OperatorDef()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+
+    input_name = get_input_tensor(bn_ops[3], 0).name
+    gamma = get_input_tensor(bn_ops[2], 1).name
+    beta = get_input_tensor(bn_ops[5], 0).name
+    mean = get_input_tensor(bn_ops[4], 0).name
+    variance = get_input_tensor(bn_ops[0], 0).name
+
+    op_def.name = op.name[:-4]  # remove /add
+    op_def.type = 'BatchNorm'
+    if self.device == 'gpu':
+      op_def.input.extend([input_name])
+      for tensor_name in [gamma, beta, mean, variance]:
+        output_name = self.add_buffer_to_image(tensor_name, "ARGUMENT")
         op_def.input.extend([output_name])
-      else:
-        op_def.input.extend([first_op.inputs[1].name])
-      op_def.output.extend([output.name for output in first_op.outputs])
-      add_output_shape(first_op.outputs, op_def)
-    elif first_op.type in ['Relu', 'SpaceToBatchND', 'BatchToSpaceND']:
-      op_def.name = first_op.name
-      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])
-      add_output_shape(first_op.outputs, op_def)
     else:
-      raise Exception('Unknown Op: %s, type: %s' % (first_op.name, first_op.type))
-      pass
+      op_def.input.extend([input_name, gamma, beta, mean, variance])
+    op_def.output.extend([output.name for output in bn_ops[6].outputs])
+    self.add_output_shape(bn_ops[6].outputs, op_def)
+    epsilon_arg = op_def.arg.add()
+    epsilon_arg.name = 'epsilon'
+    epsilon_arg.f = get_input_tensor(op, 1).eval().astype(np.float)
+    data_format_arg = op_def.arg.add()
+    data_format_arg.name = 'data_format'
+    data_format_arg.s = 'NHWC'
+
+    self.net_def.op.extend([op_def])
+    for i in range(1, 7):
+      self.resolved_ops[bn_ops[i].name] = 1
+
+  def convert_pooling(self, op):
+    op_def = self.net_def.op.add()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = 'Pooling'
+    op_def.input.extend([input.name for input in op.inputs])
+    op_def.output.extend([output.name for output in op.outputs])
+    self.add_output_shape(op.outputs, op_def)
+    pooling_type_arg = op_def.arg.add()
+    pooling_type_arg.name = 'pooling_type'
+    pooling_type_arg.i = pooling_type_mode[op.type]
+    padding_arg = op_def.arg.add()
+    padding_arg.name = 'padding'
+    padding_arg.i = padding_mode[op.get_attr('padding')]
+    strides_arg = op_def.arg.add()
+    strides_arg.name = 'strides'
+    strides_arg.ints.extend(op.get_attr('strides')[1:3])
+    kernels_arg = op_def.arg.add()
+    kernels_arg.name = 'kernels'
+    kernels_arg.ints.extend(op.get_attr('ksize')[1:3])
+    data_format_arg = op_def.arg.add()
+    data_format_arg.name = 'data_format'
+    data_format_arg.s = 'NHWC'
+    self.resolved_ops[op.name] = 1
+
+  def convert_relu6(self, op):
+    op_def = self.net_def.op.add()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = 'Relu'
+    op_def.input.extend([input.name for input in op.inputs])
+    op_def.output.extend([output.name for output in op.outputs])
+    self.add_output_shape(op.outputs, op_def)
+    max_limit_arg = op_def.arg.add()
+    max_limit_arg.name = 'max_limit'
+    max_limit_arg.f = 6
+    self.resolved_ops[op.name] = 1
+
+  def convert_add(self, op):
+    op_def = self.net_def.op.add()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = "AddN"
+    op_def.input.extend([input.name for input in op.inputs])
+    op_def.output.extend([output.name for output in op.outputs])
+    self.add_output_shape(op.outputs, op_def)
+    self.resolved_ops[op.name] = 1
+
+  def convert_concat(self, op):
+    op_def = self.net_def.op.add()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = "Concat"
+    op_def.input.extend([op.inputs[i].name for i in xrange(2)])
+    op_def.output.extend([output.name for output in op.outputs])
+    axis_arg = op_def.arg.add()
+    axis_arg.name = 'axis'
+    axis_arg.i = get_input_tensor(op, 2).eval().astype(np.int32)
+    self.add_output_shape(op.outputs, op_def)
+    self.resolved_ops[op.name] = 1
+
+  def convert_resize_bilinear(self, op):
+    op_def = self.net_def.op.add()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = "ResizeBilinear"
+    op_def.input.extend([op.inputs[0].name])
+    op_def.output.extend([output.name for output in op.outputs])
+    size_arg = op_def.arg.add()
+    size_arg.name = 'size'
+    size_arg.ints.extend(get_input_tensor(op, 1).eval().astype(np.int32).flat)
+    size_arg = op_def.arg.add()
+    size_arg.name = 'align_corners'
+    size_arg.i = op.get_attr('align_corners')
+    self.add_output_shape(op.outputs, op_def)
+    self.resolved_ops[op.name] = 1
+
+  def convert_bias_add(self, op):
+    op_def = mace_pb2.OperatorDef()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = "BiasAdd"
+    op_def.input.extend([op.inputs[0].name])
+    if self.device == 'gpu':
+      output_name = self.add_buffer_to_image(op.inputs[1].name, "ARGUMENT")
+      op_def.input.extend([output_name])
+    else:
+      op_def.input.extend([op.inputs[1].name])
+    op_def.output.extend([output.name for output in op.outputs])
+    self.add_output_shape(op.outputs, op_def)
+    self.net_def.op.extend([op_def])
+    self.resolved_ops[op.name] = 1
+
+  def convert_normal_op(self, op):
+    op_def = self.net_def.op.add()
+    arg = op_def.arg.add()
+    arg.name = 'T'
+    arg.i = self.dt
+    op_def.name = op.name
+    op_def.type = op.type
+    op_def.input.extend([input.name for input in op.inputs])
+    op_def.output.extend([output.name for output in op.outputs])
+    self.add_output_shape(op.outputs, op_def)
+    self.resolved_ops[op.name] = 1
+
+  def convert(self, input_node, output_node):
+    if self.device == 'gpu':
+      self.add_input_transform(input_node)
+
+    for op in self.tf_ops:
+      if self.resolved_ops[op.name] == 1:
+        continue
+      if op.type in ['Placeholder', 'Reshape', 'Identity']:
+        self.resolved_ops[op.name] = 1
+        pass
+      elif op.type == 'Const':
+        self.convert_tensor(op)
+      elif op.type == 'Conv2D' or op.type == 'DepthwiseConv2dNative':
+        self.convert_conv2d(op)
+      elif op.type == 'FusedBatchNorm':
+        self.convert_fused_batchnorm(op)
+      elif op.type == 'Add' and op.name.endswith('batchnorm/add'):
+        self.convert_batchnorm(op)
+      elif op.type == 'AvgPool' or op.type == 'MaxPool':
+        self.convert_pooling(op)
+      elif op.type == 'Relu6':
+        self.convert_relu6(op)
+      elif op.type == 'Add':
+        self.convert_add(op)
+      elif op.type == 'ConcatV2':
+        self.convert_concat(op)
+      elif op.type == 'ResizeBilinear':
+        self.convert_resize_bilinear(op)
+      elif op.type == 'BiasAdd':
+        self.convert_bias_add(op)
+      elif op.type in ['Relu', 'SpaceToBatchND', 'BatchToSpaceND']:
+        self.convert_normal_op(op)
+      else:
+        raise Exception('Unknown Op: %s, type: %s' % (op.name, op.type))
 
-  for i in range(resolved_count):
-    del unresolved_ops[0]
+    if self.device == 'gpu':
+      self.add_output_transform(output_node)
 
+    for key in self.resolved_ops:
+      if self.resolved_ops[key] != 1:
+        print 'Unresolve Op: %s' % key
 
 def convert_to_mace_pb(input_graph_def, input_node, output_node, data_type, device):
   net_def = mace_pb2.NetDef()
@@ -331,13 +416,8 @@ def convert_to_mace_pb(input_graph_def, input_node, output_node, data_type, devi
     with session.graph.as_default() as graph:
       tf.import_graph_def(input_graph_def, name="")
       ops = graph.get_operations()
-      unresolved_ops = ops
-      if device == 'gpu':
-        add_input_transform(input_node, dt, net_def)
-      while len(unresolved_ops) > 0:
-        convert_ops(unresolved_ops, dt, net_def, device)
-      if device == 'gpu':
-        add_output_transform(output_node, net_def)
+      converter = TFConverter(ops, net_def, dt, device)
+      converter.convert(input_node, output_node)
 
   print "PB Parsed."
 
-- 
GitLab