Operator selector: slim operator table (#1309)

Modify operator table from source model file
Generated model file's operator table has only used operators

Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>

diff --git a/tools/tflitefile_tool/select_operator.py b/tools/tflitefile_tool/select_operator.py
index ad8fa7f..db47c76 100755 (executable)
--- a/tools/tflitefile_tool/select_operator.py
+++ b/tools/tflitefile_tool/select_operator.py
@@ -40,7 +40,7 @@ def GetOperatorList(oplist_file):
     return opcode_list
 
 
-def GenerateOperatorCodes(new_builder, sample_model):
+def GenerateOperatorCodes(new_builder, sample_model, used_operators_dic):
     operator_code_num = sample_model.OperatorCodesLength()
     new_operator_code_list = []
     new_operator_code_string_list = {}
@@ -50,35 +50,38 @@ def GenerateOperatorCodes(new_builder, sample_model):
 
     # Create operator_code string
     for operator_code_idx in range(operator_code_num):
-        operator_code = sample_model.OperatorCodes(operator_code_idx)
-        operator_code_string = operator_code.CustomCode()
-        if (operator_code_string !=
-                "") and (not operator_code_string in new_operator_code_string_list):
-            new_operator_code_string_list[
-                operator_code_string] = new_builder.CreateString(operator_code_string)
+        if operator_code_idx in used_operators_dic:
+            operator_code = sample_model.OperatorCodes(operator_code_idx)
+            operator_code_string = operator_code.CustomCode()
+            if (operator_code_string !=
+                    "") and (not operator_code_string in new_operator_code_string_list):
+                new_operator_code_string_list[
+                    operator_code_string] = new_builder.CreateString(operator_code_string)
 
     # Create tables of operator_code
     for operator_code_idx in range(operator_code_num):
-        operator_code = sample_model.OperatorCodes(operator_code_idx)
+        if operator_code_idx in used_operators_dic:
+            operator_code = sample_model.OperatorCodes(operator_code_idx)
 
-        # Create operator_code table
-        tflite.OperatorCode.OperatorCodeStart(new_builder)
-        tflite.OperatorCode.OperatorCodeAddBuiltinCode(new_builder,
-                                                       operator_code.BuiltinCode())
+            # Create operator_code table
+            tflite.OperatorCode.OperatorCodeStart(new_builder)
+            tflite.OperatorCode.OperatorCodeAddBuiltinCode(new_builder,
+                                                           operator_code.BuiltinCode())
 
-        new_operator_code_string = operator_code.CustomCode()
-        if new_operator_code_string in new_operator_code_string_list:
-            tflite.OperatorCode.OperatorCodeAddCustomCode(
-                new_builder, new_operator_code_string_list[new_operator_code_string])
-        new_operator_code = tflite.OperatorCode.OperatorCodeEnd(new_builder)
-        new_operator_code_list.append(new_operator_code)
+            new_operator_code_string = operator_code.CustomCode()
+            if new_operator_code_string in new_operator_code_string_list:
+                tflite.OperatorCode.OperatorCodeAddCustomCode(
+                    new_builder, new_operator_code_string_list[new_operator_code_string])
+            new_operator_code = tflite.OperatorCode.OperatorCodeEnd(new_builder)
+            new_operator_code_list.append(new_operator_code)
 
     # Create operator_code vector
-    tflite.Model.ModelStartOperatorCodesVector(new_builder, operator_code_num)
-    for operator_code_idx in reversed(range(operator_code_num)):
+    new_operator_code_num = len(new_operator_code_list)
+    tflite.Model.ModelStartOperatorCodesVector(new_builder, new_operator_code_num)
+    for operator_code_idx in reversed(range(new_operator_code_num)):
         new_builder.PrependUOffsetTRelative(new_operator_code_list[operator_code_idx])
 
-    return new_builder.EndVector(operator_code_num)
+    return new_builder.EndVector(new_operator_code_num)
 
 
 def GenerateQuantization(new_builder, selected_quantization):
@@ -404,10 +407,12 @@ def GenerateBuiltinOption(new_builder, selected_builtin_option, builtin_option_t
     return 0
 
 
-def GenerateOperator(new_builder, selected_operator, used_tensors_dic):
+def GenerateOperator(new_builder, selected_operator, used_tensors_dic,
+                     used_operators_dic):
 
     # define opcode_index
     opcode_index = selected_operator.OpcodeIndex()
+    new_opcode_index = used_operators_dic[opcode_index]
 
     # create input vector
     input_num = selected_operator.InputsLength()
@@ -451,7 +456,7 @@ def GenerateOperator(new_builder, selected_operator, used_tensors_dic):
 
     # Create operator
     tflite.Operator.OperatorStart(new_builder)
-    tflite.Operator.OperatorAddOpcodeIndex(new_builder, opcode_index)
+    tflite.Operator.OperatorAddOpcodeIndex(new_builder, new_opcode_index)
     if input_num != 0:
         tflite.Operator.OperatorAddInputs(new_builder, new_input)
     if output_num != 0:
@@ -465,7 +470,8 @@ def GenerateOperator(new_builder, selected_operator, used_tensors_dic):
     return tflite.Operator.OperatorEnd(new_builder)
 
 
-def GenerateOperators(new_builder, selected_subgraph, opcode_list, used_tensors_dic):
+def GenerateOperators(new_builder, selected_subgraph, opcode_list, used_tensors_dic,
+                      used_operators_dic):
     operator_num = selected_subgraph.OperatorsLength()
     new_operator_list = []
 
@@ -476,7 +482,7 @@ def GenerateOperators(new_builder, selected_subgraph, opcode_list, used_tensors_
         if operator_idx in opcode_list:
             selected_operator = selected_subgraph.Operators(operator_idx)
             new_operator = GenerateOperator(new_builder, selected_operator,
-                                            used_tensors_dic)
+                                            used_tensors_dic, used_operators_dic)
             new_operator_list.append(new_operator)
 
     new_operator_num = len(new_operator_list)
@@ -491,7 +497,8 @@ def GenerateOperators(new_builder, selected_subgraph, opcode_list, used_tensors_
 
 
 def GenerateSubgraph(new_builder, selected_subgraph, opcode_list, new_input_tensor,
-                     new_output_tensor, used_tensors_dic, used_buffers_dic):
+                     new_output_tensor, used_tensors_dic, used_buffers_dic,
+                     used_operators_dic):
 
     # Tensors
     tensors = GenerateTensors(new_builder, selected_subgraph, used_tensors_dic,
@@ -517,7 +524,7 @@ def GenerateSubgraph(new_builder, selected_subgraph, opcode_list, new_input_tens
 
     # Operators
     operators = GenerateOperators(new_builder, selected_subgraph, opcode_list,
-                                  used_tensors_dic)
+                                  used_tensors_dic, used_operators_dic)
 
     # Name
     subgraph_name = selected_subgraph.Name()
@@ -540,14 +547,15 @@ def GenerateSubgraph(new_builder, selected_subgraph, opcode_list, new_input_tens
 
 
 def GenerateSubgraphs(new_builder, sample_model, opcode_list, new_input_tensor,
-                      new_output_tensor, used_tensors_dic, used_buffers_dic):
+                      new_output_tensor, used_tensors_dic, used_buffers_dic,
+                      used_operators_dic):
     new_subgraph_list = []
 
     # We think only main graph
     selected_subgraph = sample_model.Subgraphs(0)
     new_subgraph = GenerateSubgraph(new_builder, selected_subgraph, opcode_list,
                                     new_input_tensor, new_output_tensor, used_tensors_dic,
-                                    used_buffers_dic)
+                                    used_buffers_dic, used_operators_dic)
     new_subgraph_list.append(new_subgraph)
 
     new_subgraph_num = 1
@@ -603,17 +611,18 @@ def GenerateBuffers(new_builder, sample_model, used_buffers_dic):
 
 
 def GenerateModel(new_builder, sample_model, opcode_list, new_input_tensors,
-                  new_output_tensors, used_tensors_dic, used_buffers_dic):
+                  new_output_tensors, used_tensors_dic, used_buffers_dic,
+                  used_operators_dic):
     # uint
     version = sample_model.Version()
 
     # pointer of operator code 'table' vector
-    operator_codes = GenerateOperatorCodes(new_builder, sample_model)
+    operator_codes = GenerateOperatorCodes(new_builder, sample_model, used_operators_dic)
 
     # subgraphs
     subgraphs = GenerateSubgraphs(new_builder, sample_model, opcode_list,
                                   new_input_tensors, new_output_tensors, used_tensors_dic,
-                                  used_buffers_dic)
+                                  used_buffers_dic, used_operators_dic)
 
     # description
     description_string = new_builder.CreateString(sample_model.Description())
@@ -664,8 +673,9 @@ def main(args):
     sample_model = tflite.Model.Model.GetRootAsModel(sample_buf, 0)
     sample_subgraph = sample_model.Subgraphs(0)
 
-    # Collect used tensor
+    # Collect used tensor & used operator
     used_tensors = []
+    used_operators = []
 
     for opcode_idx in opcode_list:
         opcode = sample_subgraph.Operators(opcode_idx)
@@ -681,7 +691,12 @@ def main(args):
                 # default: same as input sample
                 used_tensors.append(output_tensor_idx)
 
+        opcode_idx = opcode.OpcodeIndex()
+        if not opcode_idx in used_operators:
+            used_operators.append(opcode_idx)
+
     used_tensors.sort()
+    used_operators.sort()
 
     # Collect used buffer
     # buffer[0] should be blank. So it should start from 1
@@ -693,6 +708,13 @@ def main(args):
         used_buffers.append(buf_idx)
     used_buffers.sort()
 
+    # Assign new index for operator
+    used_operators_dic = {}
+
+    for new_operator_idx in range(len(used_operators)):
+        sample_operator_idx = used_operators[new_operator_idx]
+        used_operators_dic[sample_operator_idx] = new_operator_idx
+
     # Assign new index for tensor
     used_tensors_dic = {}
 
@@ -748,7 +770,8 @@ def main(args):
     new_builder = flatbuffers.Builder(1024)
 
     new_model = GenerateModel(new_builder, sample_model, opcode_list, new_input_tensors,
-                              new_output_tensors, used_tensors_dic, used_buffers_dic)
+                              new_output_tensors, used_tensors_dic, used_buffers_dic,
+                              used_operators_dic)
 
     Finish(new_builder, new_model)
     new_buf = new_builder.Output()