updated this branch to use the new ugraph-constructor

neil-tan · neil-tan · commit f4b87e5b1bf4 · 2019-12-12T18:33:47.000+08:00
diff --git a/tests/tflm/tflite_export/conftest.py b/tests/tflm/tflite_export/conftest.py
@@ -3,148 +3,137 @@
 
 import tensorflow as tf
 from utensor_cgen.ir import TensorInfo, OperationInfo, uTensorGraph
+from utensor_cgen.ir.converter import (AttrValueConverter, DataTypeConverter,
+                                       GenericTensorConverterMixin)
 from utensor_cgen.utils import prune_graph, topologic_order_graph
-
-@fixture(name='hybrid_quant_output')
-def simple_tflm_graph():
-    ugraph = uTensorGraph()
-
-    mock_input_op0 = OperationInfo(
-        name = "mock_input_const0",
-        op_type = "Const",
-        backend = "tensorflow",
-        ugraph = ugraph,
-        op_attr = dict(),
-        input_tensors = [],
-        output_tensors = []
-    )
-    mock_input_op0.op_attr["value"] = np.array([[2],[4],[6],[8]], dtype=np.float32)
-    mock_input_op0.op_attr["shape"] = [4,1]
-
-    input0 = TensorInfo(
-        name = "input0",
-        op_name = "mock_input_const0",
-        dtype = mock_input_op0.op_attr["value"].dtype,
-        shape = mock_input_op0.op_attr["shape"],
-        ugraph = ugraph
-    )
-
-    mock_input_op0.output_tensors = [input0]
-
-    mock_input1_op = OperationInfo(
-        name = "mock_input_const1",
-        op_type = "Const",
-        backend = "tensorflow",
-        ugraph = ugraph,
-        op_attr = dict(),
-        input_tensors = [],
-        output_tensors = []
-    )
-    mock_input1_op.op_attr["value"] = np.array([[2],[4],[6],[8]], dtype=np.float32)
-    mock_input1_op.op_attr["shape"] = [4,1]
-
-    input1 = TensorInfo(
-        name = "input1",
-        op_name = "mock_input_const1",
-        dtype = mock_input1_op.op_attr["value"].dtype,
-        shape = mock_input1_op.op_attr["shape"],
-        ugraph = ugraph
-    )
-
-    mock_input1_op.output_tensors = [input1]
-
-    add_output = TensorInfo(
-        name = "add_out",
-        op_name = "add0",
-        dtype = mock_input_op0.op_attr["value"].dtype,
-        shape = mock_input_op0.op_attr["shape"],
-        ugraph = ugraph
-    )
-
-    add_op = OperationInfo(
-        name = "add0",
-        op_type = "ADD",
-        backend = "tensorflow",
-        ugraph = ugraph,
-        op_attr = dict(),
-        input_tensors = [input0, input1],
-        output_tensors = [add_output]
+from utensor_cgen.backend.operators import OperatorFactory, _Operator
+from utensor_cgen.matcher import OpEqualityDelegate, _morphism
+
+
+@OperatorFactory.register
+@OpEqualityDelegate.is_associative(
+  permutations=((0, 1), (1, 0))
+)
+class _TFLM_AddOperator(_Operator):
+
+  op_type = "TFLM_ADD" # tf op type
+
+  def __init__(self, op_info, **kwargs):
+    _Operator.__init__(self)
+    inputs = [tensor_info.name for tensor_info in op_info.input_tensors]
+    output = op_info.output_tensors[0].name
+    tf_dtype = op_info.input_tensors[0].dtype
+
+  @classmethod
+  def build_op_info(cls, ugraph, name, tensor_x, tensor_y, **kwargs):
+    # broadcast the shape and promote types
+    dummy_x = np.empty(tensor_x.shape)
+    dummy_y = np.empty(tensor_y.shape)
+    output_shape = np.broadcast(dummy_x, dummy_y).shape
+    output_dtype = np.promote_types(tensor_x.dtype, tensor_y.dtype)
+    return OperationInfo(
+      name=name,
+      input_tensors=[tensor_x, tensor_y],
+      output_tensors=[
+        TensorInfo(
+          name='{}:0'.format(name),
+          op_name=name,
+          dtype=output_dtype,
+          shape=list(output_shape),
+          ugraph=ugraph
+        )
+      ],
+      op_type=cls.op_type,
+      op_attr={
+        'T': AttrValueConverter.__utensor_generic_type__(
+          value_name='type',
+          value=DataTypeConverter.get_tf_value(output_dtype)
+        )
+      },
+      ugraph=ugraph,
+      backend=kwargs.get('backend', 'TFLM')
     )
 
-    ugraph.ops_info["ADD0"] = add_op
 
-    weight_op = OperationInfo(
-        name = "weight_const",
-        op_type = "Const",
-        backend = "tensorflow",
-        ugraph = ugraph,
-        op_attr = dict(),
-        input_tensors = [],
-        output_tensors = []
+@OperatorFactory.register
+class _TFLM_FULLY_CONNECTED_Operator(_Operator):
+
+  op_type="TFLM_FULLY_CONNECTED"
+  
+  def __init__(self, op_info, **kwargs):
+    _Operator.__init__(self)
+    inputs = [tensor_info.name for tensor_info in op_info.input_tensors]
+    output = op_info.output_tensors[0].name
+    out_dtype = op_info.output_tensors[0].dtype
+    in_dtypes = [tensor_info.dtype for tensor_info in op_info.input_tensors]
+    #assert (op_info.input_tensors[0].shape[1] == None or op_info.input_tensors[0].shape[1] == 1)
+
+  @classmethod
+  def build_op_info(cls, ugraph, name, tensor_x, tensor_w, tensor_b, **kwargs):
+    output_shape = [tensor_w.shape[0], tensor_x.shape[1]]
+    #output_dtype = np.promote_types(tensor_x.dtype, tensor_y.dtype)
+    output_dtype = tensor_x.dtype
+    return OperationInfo(
+      name=name,
+      input_tensors=[tensor_x, tensor_w, tensor_b],
+      output_tensors=[
+        TensorInfo(
+          name='{}:0'.format(name),
+          op_name=name,
+          dtype=output_dtype,
+          shape=list(output_shape),
+          ugraph=ugraph
+        )
+      ],
+      op_type=cls.op_type,
+      op_attr={
+        'T': AttrValueConverter.__utensor_generic_type__(
+          value_name='type',
+          value=DataTypeConverter.get_tf_value(output_dtype)
+        )
+      },
+      ugraph=ugraph,
+      backend=kwargs.get('backend', 'TFLM')
     )
-    #weight_op.op_attr["value"] = np.array([1,2,3,4], dtype=np.int8)
-    weight_op.op_attr["value"] = np.array([10,20,30,40], dtype=np.float32)
-    weight_op.op_attr["shape"] = [1,4]
 
-    weight = TensorInfo(
-        name = "weight",
-        op_name = "weight_const",
-        dtype = np.dtype("float32"),
-        shape = weight_op.op_attr["shape"],
-        ugraph = ugraph
-    )
-    weight_op.output_tensors = [weight]
-
-    bias_op = OperationInfo(
-        name = "bias_const",
-        op_type = "Const",
-        backend = "tensorflow",
-        ugraph = ugraph,
-        op_attr = dict(),
-        input_tensors = [],
-        output_tensors = []
-    )
-    #bias_op.op_attr["value"] = np.array([1], dtype=np.int8)
-    bias_op.op_attr["value"] = np.array([7], dtype=np.float32)
-    bias_op.op_attr["shape"] = [1]
-
-    bias = TensorInfo(
-        name = "bias",
-        op_name = "bias_const",
-        dtype = np.dtype("float32"),
-        shape = bias_op.op_attr["shape"],
-        ugraph = ugraph
-    )
-    bias_op.output_tensors = [bias]
-
-
-    fc1_op = OperationInfo(
-        name = "FC1",
-        op_type = "FULLY_CONNECTED",
-        backend = "tensorflow",
-        ugraph = ugraph,
-        op_attr = dict(),
-        input_tensors = [],
-        output_tensors = []
-    )
-
-    output = TensorInfo(
-        name = "output",
-        op_name = "FC1",
-        dtype = np.dtype("float32"),
-        shape = [1],
-        ugraph = ugraph
-    )
-
-    fc1_op.input_tensors = [add_output, weight, bias]
-    fc1_op.output_tensors = [output]
-
-    ugraph.ops_info["FC1"] = fc1_op
-    ugraph.output_nodes = ["FC1"]
-    #ugraph.backend = "tensorflow"
-
-    topologic_order_graph(ugraph)
-    #ugraph = prune_graph(ugraph)
+@fixture(name='hybrid_quant_output')
+def simple_tflm_graph():
+    ugraph = uTensorGraph()
 
-    #return: ugraph, input tensors, output tensors
-    return [ugraph, ["input0", "input1"], ["weight", "bias", "output"]]
+    with ugraph.begin_construction():
+        tensor_x0, = ugraph.add_op(
+            op_type='Const',
+            name='x0',
+            value=np.array([1, 1, 1, 1], dtype=np.float32)[:, np.newaxis]
+        )
+        tensor_x1, = ugraph.add_op(
+            op_type='Const',
+            name='x1',
+            value=np.array([2, 4, 6, 8], dtype=np.float32)[:, np.newaxis]
+        )
+        tensor_w, = ugraph.add_op(
+            op_type='Const',
+            name='w',
+            value=np.array([10, 20, 30, 40], dtype=np.float32)[np.newaxis, :]
+        )
+        tensor_b, = ugraph.add_op(
+            op_type='Const',
+            name='b',
+            value=np.array([7], dtype=np.float32)
+        )
+        
+
+        tensor_addout, = ugraph.add_op(
+            tensor_x0, tensor_x1,
+            op_type='TFLM_ADD',
+            name='TFLM_ADD0'
+        )
+
+        tensor_out, = ugraph.add_op(
+            tensor_addout, tensor_w, tensor_b,
+            op_type='TFLM_FULLY_CONNECTED',
+            name='TFLM_FULLY_CONNECTED00',
+            is_output=True
+        )
+
+    return [ugraph, ["x0:0", "x1:0"], ["w:0", "b:0", tensor_out.name]]
diff --git a/utensor_cgen/transformer/tflite_exporter.py b/utensor_cgen/transformer/tflite_exporter.py
@@ -23,6 +23,42 @@
 
 __all__ = ["TFLiteExporter"]
 
+def rename_ugraph_ops_(ugraph, name_old, name_new):
+  op = ugraph.ops_info.pop(name_old)
+  op.name = name_new
+  for out_tensor in op.output_tensors:
+    pass # renaming tensors
+  for out_node in op.output_nodes:
+    for in_tensor in out_node.input_tensors:
+      if in_tensor.op_name == name_old:
+        in_tensor.op_name = name_new
+  ugraph.ops_info[name_new] = op
+  for i, op_name in enumerate(ugraph.topo_order):
+    if op_name == name_old:
+      ugraph.topo_order[i] = name_new
+      break
+  for i, op_name in enumerate(ugraph.output_nodes):
+    if op_name == name_old:
+      ugraph.output_nodes[i] = name_new
+      break
+
+  op.name = name_new
+  for out_tensor in op.output_tensors:
+    pass # renaming tensors
+  for out_node in op.output_nodes:
+    for in_tensor in out_node.input_tensors:
+      if in_tensor.op_name == name_old:
+        in_tensor.op_name = name_new
+  ugraph.ops_info[name_new] = op
+  for i, op_name in enumerate(ugraph.topo_order):
+    if op_name == name_old:
+      ugraph.topo_order[i] = name_new
+      break
+  for i, op_name in enumerate(ugraph.output_nodes):
+    if op_name == name_old:
+      ugraph.output_nodes[i] = name_new
+      break
+
 def get_fullyconnected_builtin_option(fbuilder, op_info):
 
   weight_format = tflite.FullyConnectedOptionsWeightsFormat.FullyConnectedOptionsWeightsFormat.DEFAULT
@@ -110,6 +146,7 @@ def __init__(self, input_tensors, output_tensors):
 
 
   def transform(self, ugraph):
+    self.__tflm_graph_legalize_(ugraph)
     # create tensor data buffer
     # create const tensors
     # update tensor_index
@@ -184,7 +221,7 @@ def __create_static_tensor(self, ugraph):
                         out_tensor_info.shape)
         #weight
         #value = op_info.op_attr['value'].value.np_array.flatten() #TODO: specify endianness here
-        value = op_info.op_attr['value'].flatten() #FIXME: deal with the proper attribute type here 
+        value = op_info.op_attr['value'].value.np_array.flatten() #FIXME: deal with the proper attribute type here 
         raw_data = value.tobytes()
         data_vec = self.__fb_vector(tflite.Buffer.BufferStartDataVector, raw_data)
         
@@ -284,7 +321,7 @@ def __sym_quantization_info(self, op_info):
     """
 
     #values = op_info.op_attr['value'].value.np_array.flatten()
-    values = op_info.op_attr['value'].flatten() #FIXME: deal with the proper attribute type here 
+    values = op_info.op_attr['value'].value.np_array.flatten() #FIXME: deal with the proper attribute type here 
     abs_max = np.absolute(values).max()
     #based on quantizted int8 dtype
     scale = 127 / abs_max
@@ -366,5 +403,10 @@ def __fb_vector(self, vector_constructor, arr_list):
     
     return self.fbuilder.EndVector(num_items)
 
+  def __tflm_graph_legalize_(self, ugraph):
+    for _, op in ugraph.ops_info.items():
+      assert op.op_type.startswith("TFLM_") or op.op_type in self.static_op_types
+      #TODO: strip the leading TFLM_ only
+      op.op_type = op.op_type.replace("TFLM_", '')
 # How to construct the quantization parameter for intermediate tensors?
 ## zero point and scale
