input1 tensor failed to print after invoke

Author: neil-tan

tests/tflm/tflite_export/conftest.py

@@ -18,13 +18,14 @@ def simple_tflm_graph():
         input_tensors = [],
         output_tensors = []
     )
-    weight_op.op_attr["value"] = np.array([1,2,3,4], dtype=np.int8)
-    weight_op.op_attr["shape"] = [4,1]
+    #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("int8"),
+        dtype = np.dtype("float32"),
         shape = weight_op.op_attr["shape"],
         ugraph = ugraph
     )
@@ -39,14 +40,14 @@ def simple_tflm_graph():
         input_tensors = [],
         output_tensors = []
     )
-    mock_input_op.op_attr["value"] = np.array([[1.0],[2.0],[3.0],[4.0]], dtype=np.float32)
-    mock_input_op.op_attr["shape"] = [1,4]
+    mock_input_op.op_attr["value"] = np.array([[2],[4],[6],[8]], dtype=np.float32)
+    mock_input_op.op_attr["shape"] = [4,1]
 
     input1 = TensorInfo(
-        name = "input",
+        name = "input1",
         op_name = "mock_input_const",
         dtype = mock_input_op.op_attr["value"].dtype,
-        shape = [1, 4],
+        shape = mock_input_op.op_attr["shape"],
         ugraph = ugraph
     )
 
@@ -61,13 +62,14 @@ def simple_tflm_graph():
         input_tensors = [],
         output_tensors = []
     )
-    bias_op.op_attr["value"] = np.array([1], dtype=np.int8)
+    #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("int8"),
+        dtype = np.dtype("float32"),
         shape = bias_op.op_attr["shape"],
         ugraph = ugraph
     )
@@ -87,8 +89,8 @@ def simple_tflm_graph():
     output = TensorInfo(
         name = "output",
         op_name = "FC1",
-        dtype = np.dtype("float"),
-        shape = [1, 1],
+        dtype = np.dtype("float32"),
+        shape = [1],
         ugraph = ugraph
     )
 
@@ -103,4 +105,4 @@ def simple_tflm_graph():
     #ugraph = prune_graph(ugraph)
 
     #return: ugraph, input tensors, output tensors
-    return [ugraph, [], ["output"]]
+    return [ugraph, [], ["input1", "weight", "bias", "output"]]

tests/tflm/tflite_export/test_write.py

@@ -86,19 +86,21 @@ def test_tflite_fb_write(hybrid_quant_output):
     ugraph = exporter.transform(sample_ugraph)
     model_content = exporter.output()
 
-    print_tflite_graph(model_content)
+    #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()
-    test_model_output_index = test_model.tensor(test_model.get_output_details()[0]["index"])
     test_model.invoke()
-    output_content = test_model.get_tensor(test_model_output_index)[0]
 
-    print(output_content)
+    print(test_model.get_tensor_details())
+    print("1 :", test_model.get_tensor(test_model.get_output_details()[0]["index"]))
+    print("2 :", test_model.get_tensor(test_model.get_output_details()[1]["index"]))
+    print("3 :", test_model.get_tensor(test_model.get_output_details()[2]["index"]))
+    print("out :", test_model.get_tensor(test_model.get_output_details()[3]["index"]))
 
     test_pass = True
     assert test_pass, 'error message here'

utensor_cgen/transformer/tflite_exporter.py

@@ -38,6 +38,15 @@ def get_fullyconnected_builtin_option(fbuilder, op_info):
 
   return obj, BuiltinOptions.FullyConnectedOptions
 
+def tensor_type_lookup(numpy_dtype):
+  TensorType = tflite.TensorType.TensorType
+  lookup_map = dict()
+  lookup_map[np.dtype('float32')] = TensorType.FLOAT32
+  lookup_map[np.dtype('int8')] = TensorType.INT8
+
+  return lookup_map[numpy_dtype]
+
+
 class FlatbufferOpManager:
     op_list = list()
     code_name_lookup = {v: k for k, v in BuiltinOperator.__dict__.items()}
@@ -193,7 +202,7 @@ def __create_static_tensor(self, ugraph):
         #tensor object
         tflite.Tensor.TensorStart(self.fbuilder)
         tflite.Tensor.TensorAddShape(self.fbuilder, shape_vec)
-        tflite.Tensor.TensorAddType(self.fbuilder, TensorType.INT8) #TODO: a conversion class here, out_dtype
+        tflite.Tensor.TensorAddType(self.fbuilder, tensor_type_lookup(out_dtype))
         if export_tensor_name:
           tflite.Tensor.TensorAddName(self.fbuilder, tensor_name)
         tflite.Tensor.TensorAddQuantization(self.fbuilder, q_param)
@@ -306,7 +315,7 @@ def __create_variable_tensors(self, ugraph):
       tflite.Tensor.TensorStart(self.fbuilder)
       tflite.Tensor.TensorAddShape(self.fbuilder, shape_vec)
       #tflite.Tensor.TensorAddType(self.fbuilder, TensorType.INT8) #TODO: tensor type conversion here
-      tflite.Tensor.TensorAddType(self.fbuilder, TensorType.FLOAT32)
+      tflite.Tensor.TensorAddType(self.fbuilder, tensor_type_lookup(tensor_info.dtype))
       tflite.Tensor.TensorAddName(self.fbuilder, tensor_name)
       #tflite.Tensor.TensorAddQuantization(self.fbuilder, q_param)
       tflite.Tensor.TensorAddIsVariable(self.fbuilder, True)