Merge pull request #651 from PINTO0309/sm_to_cf

Workaround to the problem of TensorFlow corrupting the FlatBuffer input/output order during INT8 quantization.

Author: PINTO0309

README.md

@@ -293,15 +293,15 @@ Video speed is adjusted approximately 50 times slower than actual speed.
   docker run --rm -it \
   -v `pwd`:/workdir \
   -w /workdir \
-  ghcr.io/pinto0309/onnx2tf:1.22.4
+  ghcr.io/pinto0309/onnx2tf:1.22.5
 
   or
 
   # Authentication is not required for pulls from Docker Hub.
   docker run --rm -it \
   -v `pwd`:/workdir \
   -w /workdir \
-  docker.io/pinto0309/onnx2tf:1.22.4
+  docker.io/pinto0309/onnx2tf:1.22.5
 
   or
 

onnx2tf/__init__.py

@@ -1,3 +1,3 @@
 from onnx2tf.onnx2tf import convert, main
 
-__version__ = '1.22.4'
+__version__ = '1.22.5'

onnx2tf/onnx2tf.py

@@ -856,6 +856,12 @@ def sanitizing(node):
         onnx_graph_output_names: List[str] = [
             outputop.name for outputop in graph.outputs
         ]
+        onnx_graph_input_shapes: List[List[int | str]] = [
+            inputop.shape for inputop in graph.inputs
+        ]
+        onnx_graph_output_shapes: List[List[int | str]] = [
+            outputop.shape for outputop in graph.outputs
+        ]
 
         # Inputs
         for graph_input in graph.inputs:
@@ -1298,6 +1304,8 @@ def sanitizing(node):
                 tflite_file_name=f'{output_file_name}_float32.tflite',
                 onnx_input_names=onnx_graph_input_names,
                 onnx_output_names=onnx_graph_output_names,
+                onnx_graph_input_shapes=onnx_graph_input_shapes,
+                onnx_graph_output_shapes=onnx_graph_output_shapes,
             )
         if output_weights:
             weights_export(
@@ -1325,6 +1333,8 @@ def sanitizing(node):
                 tflite_file_name=f'{output_file_name}_float16.tflite',
                 onnx_input_names=onnx_graph_input_names,
                 onnx_output_names=onnx_graph_output_names,
+                onnx_graph_input_shapes=onnx_graph_input_shapes,
+                onnx_graph_output_shapes=onnx_graph_output_shapes,
             )
         if output_weights:
             weights_export(
@@ -1372,6 +1382,8 @@ def sanitizing(node):
                         tflite_file_name=f'{output_file_name}_dynamic_range_quant.tflite',
                         onnx_input_names=onnx_graph_input_names,
                         onnx_output_names=onnx_graph_output_names,
+                        onnx_graph_input_shapes=onnx_graph_input_shapes,
+                        onnx_graph_output_shapes=onnx_graph_output_shapes,
                     )
                 if output_weights:
                     weights_export(
@@ -1501,6 +1513,8 @@ def representative_dataset_gen():
                         tflite_file_name=f'{output_file_name}_integer_quant.tflite',
                         onnx_input_names=onnx_graph_input_names,
                         onnx_output_names=onnx_graph_output_names,
+                        onnx_graph_input_shapes=onnx_graph_input_shapes,
+                        onnx_graph_output_shapes=onnx_graph_output_shapes,
                     )
                 if output_weights:
                     weights_export(
@@ -1536,6 +1550,8 @@ def representative_dataset_gen():
                         tflite_file_name=f'{output_file_name}_full_integer_quant.tflite',
                         onnx_input_names=onnx_graph_input_names,
                         onnx_output_names=onnx_graph_output_names,
+                        onnx_graph_input_shapes=onnx_graph_input_shapes,
+                        onnx_graph_output_shapes=onnx_graph_output_shapes,
                     )
                 if output_weights:
                     weights_export(
@@ -1572,6 +1588,8 @@ def representative_dataset_gen():
                         tflite_file_name=f'{output_file_name}_integer_quant_with_int16_act.tflite',
                         onnx_input_names=onnx_graph_input_names,
                         onnx_output_names=onnx_graph_output_names,
+                        onnx_graph_input_shapes=onnx_graph_input_shapes,
+                        onnx_graph_output_shapes=onnx_graph_output_shapes,
                     )
                 info(Color.GREEN(f'INT8 Quantization with int16 activations tflite output complete!'))
             except RuntimeError as ex:
@@ -1603,6 +1621,8 @@ def representative_dataset_gen():
                         tflite_file_name=f'{output_file_name}_full_integer_quant_with_int16_act.tflite',
                         onnx_input_names=onnx_graph_input_names,
                         onnx_output_names=onnx_graph_output_names,
+                        onnx_graph_input_shapes=onnx_graph_input_shapes,
+                        onnx_graph_output_shapes=onnx_graph_output_shapes,
                     )
                 info(Color.GREEN(f'Full INT8 Quantization with int16 activations tflite output complete!'))
             except RuntimeError as ex:

onnx2tf/utils/common_functions.py

@@ -4544,6 +4544,8 @@ def rewrite_tflite_inout_opname(
     tflite_file_name: str,
     onnx_input_names: List[str],
     onnx_output_names: List[str],
+    onnx_graph_input_shapes: List[List[int |str]],
+    onnx_graph_output_shapes: List[List[int |str]],
 ):
     """Rewrite the input/output OP name of tflite to the input/output OP name of ONNX.
     Pre-installation of flatc is required.
@@ -4561,6 +4563,12 @@ def rewrite_tflite_inout_opname(
 
     onnx_output_names: List[str]
         List of ONNX output OP names
+
+    onnx_graph_input_shapes: List[List[int |str]]
+        List of ONNX input OP shapes
+
+    onnx_graph_output_shapes: List[List[int |str]]
+        List of ONNX output OP shapes
     """
     try:
         # Check to see if flatc is installed
@@ -4611,12 +4619,50 @@ def rewrite_tflite_inout_opname(
         flat_output_nums: List[int] = flat_subgraphs.outputs
         flat_input_infos = [flat_tensors[idx] for idx in flat_input_nums]
         flat_output_infos = [flat_tensors[idx] for idx in flat_output_nums]
+
+        # Determination of the number of inputs/outputs of the same shape
+        # Correct name discrepancies based on shape if multiple inputs/outputs shapes do not overlap
+        # However, if there are inputs/outputs containing undefined dimensions,
+        # workaround is skipped because correction is not possible.
+        # https://github.com/PINTO0309/onnx2tf/issues/650
+        inputs_second_dim_elements = [
+            tuple(onnx_graph_input_shape) \
+                for onnx_graph_input_shape in onnx_graph_input_shapes
+        ]
+        inputs_has_duplicates = len(inputs_second_dim_elements) != len(set(inputs_second_dim_elements))
+        inputs_has_undefined_dim = any(isinstance(item, str) for onnx_graph_input_shape in onnx_graph_input_shapes for item in onnx_graph_input_shape)
+
+        outputs_second_dim_elements = [
+            tuple(onnx_graph_output_shape) \
+                for onnx_graph_output_shape in onnx_graph_output_shapes
+        ]
+        outputs_has_duplicates = len(outputs_second_dim_elements) != len(set(outputs_second_dim_elements))
+        outputs_has_undefined_dim = any(isinstance(item, str) for onnx_graph_output_shape in onnx_graph_output_shapes for item in onnx_graph_output_shape)
+
         # INPUT
-        for idx, flat_input_info in enumerate(flat_input_infos):
-            flat_input_info.name = onnx_input_names[idx]
+        if not inputs_has_duplicates and not inputs_has_undefined_dim:
+            for onnx_input_name, onnx_input_shape in zip(onnx_input_names, onnx_graph_input_shapes):
+                for flat_input_info in flat_input_infos:
+                    if np.prod(onnx_input_shape) == np.prod(list(flat_input_info.shape)):
+                        flat_input_info.name = onnx_input_name
+                        break
+        else:
+            for idx, flat_input_info in enumerate(flat_input_infos):
+                flat_input_info.name = onnx_input_names[idx]
+
         # OUTPUT
-        for idx, flat_output_info in enumerate(flat_output_infos):
-            flat_output_info.name = onnx_output_names[idx]
+        if not outputs_has_duplicates and not outputs_has_undefined_dim:
+            for onnx_output_name, onnx_output_shape in zip(onnx_output_names, onnx_graph_output_shapes):
+                for flat_output_info in flat_output_infos:
+                    if np.prod(onnx_output_shape) == np.prod(list(flat_output_info.shape)):
+                        flat_output_info.name = onnx_output_name
+                        break
+        else:
+            for idx, flat_output_info in enumerate(flat_output_infos):
+                flat_output_info.name = onnx_output_names[idx]
+
+        if inputs_has_duplicates or inputs_has_undefined_dim or outputs_has_duplicates or outputs_has_undefined_dim:
+            warn('Carefully check the output .tflite as the order of input OP names and output OP names may have been corrupted by TensorFlow.')
 
         # make signature_defs
         """