Tflite backend

setup.py

@@ -37,7 +37,8 @@
         'torch',
         'torchvision',
         'onnx-tf==1.2.1',
-        'graphviz'
+        'graphviz',
+        'flatbuffers'
     ],
     extras_require={
         'dev': ['pytest']

tests/tflm/tflite_export/conftest.py

@@ -0,0 +1,139 @@
+import numpy as np
+from pytest import fixture
+
+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
+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')
+    )
+
+
+@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')
+    )
+
+@fixture(name='hybrid_quant_output')
+def simple_tflm_graph():
+    ugraph = uTensorGraph()
+
+    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]]

tests/tflm/tflite_export/test_write.py

@@ -0,0 +1,114 @@
+import numpy as np
+
+import tensorflow as tf
+from utensor_cgen.frontend.tensorflow import GraphDefParser
+from utensor_cgen.matcher import uTensorGraphMatcher
+from utensor_cgen.utils import prune_graph, topologic_order_graph
+from utensor_cgen.transformer import TFLiteExporter
+import utensor_cgen.third_party.flatbuffers as flatbuffers
+import utensor_cgen.third_party.tflite as tflite
+from utensor_cgen.third_party.tflite.BuiltinOperator import BuiltinOperator
+from utensor_cgen.third_party.tflite.Model import Model
+from utensor_cgen.third_party.tflite.BuiltinOptions import BuiltinOptions
+from utensor_cgen.third_party.tflite.TensorType import TensorType
+
+builtin_ops = {v: k for k, v in BuiltinOperator.__dict__.items()}
+op_options = {v: k for k, v in BuiltinOptions.__dict__.items()}
+
+tensor_np_type = dict()
+tensor_np_type[0] = np.float32
+tensor_np_type[1] = np.float16
+tensor_np_type[2] = np.int32
+tensor_np_type[3] = np.uint8
+tensor_np_type[4] = np.uint64
+tensor_np_type[5] = np.ubyte #FIXME: supposed to be string
+tensor_np_type[6] = np.bool
+tensor_np_type[7] = np.int16
+tensor_np_type[8] = np.cdouble
+tensor_np_type[9] = np.int8
+
+def print_tflite_graph(byte_buff):
+
+    model = Model.GetRootAsModel(byte_buff, 0)
+    subgraphs_len = model.SubgraphsLength()
+    subgraph = model.Subgraphs(0)
+    n_ops = subgraph.OperatorsLength()
+    print("version: ", model.Version())
+    print("subgraph len: ", subgraphs_len)
+    print("number of operators: ", n_ops)
+    print("number of t buff: ", model.BuffersLength())
+    print("flat buffer length: ", len(byte_buff), " bytes")
+    op_codes = []
+    for i in range(0, model.OperatorCodesLength()):
+        op_code =  model.OperatorCodes(i)
+        op_codes.append(op_code)
+    print("op code length: ", len(op_codes))
+
+    for i in range(0, subgraph.OperatorsLength()):
+        op = subgraph.Operators(i)
+        print("op code index: ", op.OpcodeIndex())
+        opIndex = op.OpcodeIndex()
+        op_code = op_codes[opIndex]
+        builtin_code = op_code.BuiltinCode()
+        op_type = builtin_ops[builtin_code]
+        print(op_type)
+
+        input_tensors = [subgraph.Tensors(input_idx) for input_idx in op.InputsAsNumpy()]
+        for tensor in input_tensors:
+            print()
+            print(tensor.Name(), ", ", tensor.ShapeAsNumpy())
+            print("variable: ", tensor.IsVariable())
+            if tensor.Type() == np.uint8 or tensor.Type() == np.int8:
+                q = tensor.Quantization()
+                assert q != None
+                print("quantization info: ")
+                print(" Detail Type: ", q.DetailsType())
+                print("      Scales: ", q.ScaleAsNumpy())
+                print("       Zeros: ", q.ZeroPointAsNumpy())
+                print("       Scale: ", q.ScaleAsNumpy())
+                print("  Zero Point: ", q.ZeroPointAsNumpy())      
+                print("   Dimension: ", q.QuantizedDimension())      
+
+            print(tensor.IsVariable())
+            if not tensor.IsVariable():
+                buffer_index = tensor.Buffer()
+                assert buffer_index >= 0
+                assert model.Buffers(buffer_index).DataLength() > 0
+                buffer_content = model.Buffers(buffer_index).DataAsNumpy()
+                print("Tensor values: ", buffer_content.astype(tensor_np_type[tensor.Type()]))
+            else:
+                print("None")
+
+def test_tflite_fb_write(hybrid_quant_output):
+    [sample_ugraph, input_tensors, output_tensors] = hybrid_quant_output
+    exporter = TFLiteExporter(input_tensors=input_tensors, output_tensors=output_tensors)
+    ugraph = exporter.transform(sample_ugraph)
+    model_content = exporter.output()
+
+    print_tflite_graph(model_content)
+
+    # referece_model_content = open('/Users/neitan01/Documents/tflm/sinExample/sine_model.tflite', "rb").read()
+    # print_tflite_graph(referece_model_content)
+
+    open("tflm_test_model.tflite", "wb").write(model_content)
+    test_model = tf.lite.Interpreter('tflm_test_model.tflite')
+    test_model.allocate_tensors()
+    input_data = np.array(np.ones([4,1]), dtype=np.float32)
+    test_model.set_tensor(test_model.get_input_details()[0]['index'], input_data)
+    test_model.invoke()
+
+    print("0 :", test_model.get_tensor(0))
+    print("1 :", test_model.get_tensor(1))
+    print("2 :", test_model.get_tensor(2))
+    print("3 :", test_model.get_tensor(3))
+
+
+    print(test_model.get_tensor_details())
+    print("out0 :", test_model.get_tensor(test_model.get_output_details()[0]["index"]))
+    print("out1 :", test_model.get_tensor(test_model.get_output_details()[1]["index"]))
+    print("out2 :", test_model.get_tensor(test_model.get_output_details()[2]["index"]))
+
+
+    output = test_model.get_tensor(test_model.get_output_details()[2]["index"])
+
+    assert np.abs(output - 707) <= 0.0001, 'error is greater than 0.0001'

utensor_cgen/third_party/flatbuffers/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2014 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .builder import Builder
+from .table import Table
+from .compat import range_func as compat_range