Tests for tflite postprocess (#1302)

TomWildenhain-Microsoft · web-flow · commit d7b179e9ec2e · 2021-01-29T18:59:10.000-05:00
Signed-off-by: Tom Wildenhain &lt;tomwi@microsoft.com&gt;
diff --git a/tests/backend_test_base.py b/tests/backend_test_base.py
@@ -210,11 +210,13 @@ def convert_to_tflite(self, graph_def, feed_dict, outputs):
     def run_tflite(self, tflite_path, feed_dict):
         try:
             interpreter = tf.lite.Interpreter(tflite_path)
-            interpreter.allocate_tensors()
             input_details = interpreter.get_input_details()
             output_details = interpreter.get_output_details()
             input_name_to_index = {n['name'].split(':')[0]: n['index'] for n in input_details}
             feed_dict_without_port = {k.split(':')[0]: v for k, v in feed_dict.items()}
+            for k, v in feed_dict_without_port.items():
+                interpreter.resize_tensor_input(input_name_to_index[k], v.shape)
+            interpreter.allocate_tensors()
             # The output names might be different in the tflite but the order is the same
             output_names = [n['name'] for n in output_details]
             for k, v in feed_dict_without_port.items():
diff --git a/tests/common.py b/tests/common.py
@@ -31,6 +31,7 @@
     "check_opset_max_version",
     "skip_tf2",
     "skip_tflite",
+    "requires_tflite",
     "check_opset_after_tf_version",
     "check_target",
     "skip_caffe2_backend",
@@ -195,6 +196,13 @@ def test(self):
     return decorator
 
 
+def requires_tflite(message=""):
+    """ Skip test if tflite tests are disabled """
+    config = get_test_config()
+    reason = _append_message("test requires tflite", message)
+    return unittest.skipIf(config.skip_tflite_tests, reason)
+
+
 def requires_custom_ops(message=""):
     """ Skip until custom ops framework is on PyPI. """
     reason = _append_message("test needs custom ops framework", message)
diff --git a/tests/test_tflite_postprocess.py b/tests/test_tflite_postprocess.py
@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""Unit Tests for TFLite_Detection_PostProcess op"""
+
+import os
+import struct
+import numpy as np
+import flatbuffers
+
+from common import *  # pylint: disable=wildcard-import,unused-wildcard-import
+from backend_test_base import Tf2OnnxBackendTestBase
+
+from tf2onnx import utils
+from tf2onnx.tfonnx import process_tf_graph
+from tf2onnx import optimizer
+
+from tf2onnx.tflite import Model, OperatorCode, SubGraph, Operator, Tensor, Buffer
+from tf2onnx.tflite.BuiltinOperator import BuiltinOperator
+from tf2onnx.tflite.TensorType import TensorType
+from tf2onnx.tflite.CustomOptionsFormat import CustomOptionsFormat
+
+# pylint: disable=missing-docstring
+
+
+class TFLiteDetectionPostProcessTests(Tf2OnnxBackendTestBase):
+
+    @requires_tflite("TFLite_Detection_PostProcess")
+    @check_opset_min_version(11, "Pad")
+    def test_postprocess_model1(self):
+        self._test_postprocess(num_classes=5, num_boxes=100, detections_per_class=2, max_detections=20)
+
+    @requires_tflite("TFLite_Detection_PostProcess")
+    @check_opset_min_version(11, "Pad")
+    def test_postprocess_model2(self):
+        self._test_postprocess(num_classes=5, num_boxes=100, detections_per_class=7, max_detections=20)
+
+    @requires_tflite("TFLite_Detection_PostProcess")
+    @check_opset_min_version(11, "Pad")
+    def test_postprocess_model3(self):
+        self._test_postprocess(num_classes=5, num_boxes=3, detections_per_class=7, max_detections=20)
+
+    @requires_tflite("TFLite_Detection_PostProcess")
+    @check_opset_min_version(11, "Pad")
+    def test_postprocess_model4(self):
+        self._test_postprocess(num_classes=5, num_boxes=99, detections_per_class=2, max_detections=20, extra_class=True)
+
+    def _test_postprocess(self, num_classes, num_boxes, detections_per_class, max_detections, extra_class=False):
+        model = self.make_postprocess_model(num_classes=num_classes, detections_per_class=detections_per_class,
+                                            max_detections=max_detections, x_scale=11.0, w_scale=6.0)
+
+        np.random.seed(42)
+        box_encodings_val = np.random.random_sample([1, num_boxes, 4]).astype(np.float32)
+        if extra_class:
+            num_classes += 1
+        class_predictions_val = np.random.random_sample([1, num_boxes, num_classes]).astype(np.float32)
+        anchors_val = np.random.random_sample([num_boxes, 4]).astype(np.float32)
+
+        feed_dict = {
+            "box_encodings": box_encodings_val,
+            "class_predictions": class_predictions_val,
+            "anchors": anchors_val
+        }
+
+        self.run_tflite_test(model, feed_dict)
+
+    def make_postprocess_model(self, max_detections=10, detections_per_class=100, max_classes_per_detection=1,
+                               use_regular_nms=True, nms_score_threshold=0.3, nms_iou_threshold=0.6, num_classes=90,
+                               x_scale=10.0, y_scale=10.0, w_scale=5.0, h_scale=5.0):
+        """Returns the bytes of a tflite model containing a single TFLite_Detection_PostProcess op"""
+
+        builder = flatbuffers.Builder(1024)
+
+        # op_code
+        custom_code = builder.CreateString("TFLite_Detection_PostProcess")
+        OperatorCode.OperatorCodeStart(builder)
+        OperatorCode.OperatorCodeAddDeprecatedBuiltinCode(builder, BuiltinOperator.CUSTOM)
+        OperatorCode.OperatorCodeAddCustomCode(builder, custom_code)
+        OperatorCode.OperatorCodeAddBuiltinCode(builder, BuiltinOperator.CUSTOM)
+        op_code = OperatorCode.OperatorCodeEnd(builder)
+
+        # op_codes
+        Model.ModelStartOperatorCodesVector(builder, 1)
+        builder.PrependUOffsetTRelative(op_code)
+        op_codes = builder.EndVector(1)
+
+        # Make tensors
+        # [names, shape, type tensors]
+        ts = []
+        inputs_info = [('box_encodings', [-1, -1, 4]), ('class_predictions', [-1, -1, -1]), ('anchors', [-1, 4])]
+        outputs_info = [
+            ('detection_boxes', [-1, -1, 4]),
+            ('detection_classes', [-1, -1]),
+            ('detection_scores', [-1, -1]),
+            ('num_detections', [-1])
+        ]
+        for name_info, shape_info in inputs_info + outputs_info:
+
+            name = builder.CreateString(name_info)
+            shape = builder.CreateNumpyVector(np.maximum(np.array(shape_info, np.int32), 1))
+            shape_signature = builder.CreateNumpyVector(np.array(shape_info, np.int32))
+
+            Tensor.TensorStart(builder)
+            Tensor.TensorAddShape(builder, shape)
+            Tensor.TensorAddType(builder, TensorType.FLOAT32)
+            Tensor.TensorAddName(builder, name)
+            Tensor.TensorAddShapeSignature(builder, shape_signature)
+            ts.append(Tensor.TensorEnd(builder))
+
+        SubGraph.SubGraphStartTensorsVector(builder, len(ts))
+        for tensor in reversed(ts):
+            builder.PrependUOffsetTRelative(tensor)
+        tensors = builder.EndVector(len(ts))
+
+        # inputs
+        SubGraph.SubGraphStartInputsVector(builder, 3)
+        for inp in reversed([0, 1, 2]):
+            builder.PrependInt32(inp)
+        inputs = builder.EndVector(3)
+
+        # outputs
+        SubGraph.SubGraphStartOutputsVector(builder, 4)
+        for out in reversed([3, 4, 5, 6]):
+            builder.PrependInt32(out)
+        outputs = builder.EndVector(4)
+
+        flexbuffer = \
+            b'y_scale\x00nms_score_threshold\x00max_detections\x00x_scale\x00w_scale\x00nms_iou_threshold' \
+            b'\x00use_regular_nms\x00h_scale\x00max_classes_per_detection\x00num_classes\x00detections_per_class' \
+            b'\x00\x0b\x16E>\x88j\x9e([v\x7f\xab\x0b\x00\x00\x00\x01\x00\x00\x00\x0b\x00\x00\x00*attr4**attr7*' \
+            b'*attr10**attr9**attr1**attr2**attr3**attr11*\x00\x00\x00*attr8**attr5**attr6*\x06\x0e\x06\x06\x0e' \
+            b'\x0e\x06j\x0e\x0e\x0e7&\x01'
+        flexbuffer = flexbuffer.replace(b'*attr1*', struct.pack('<f', nms_iou_threshold))
+        flexbuffer = flexbuffer.replace(b'*attr2*', struct.pack('<f', nms_score_threshold))
+        flexbuffer = flexbuffer.replace(b'*attr3*', struct.pack('<i', num_classes))
+        flexbuffer = flexbuffer.replace(b'*attr4*', struct.pack('<i', detections_per_class))
+        flexbuffer = flexbuffer.replace(b'*attr5*', struct.pack('<f', x_scale))
+        flexbuffer = flexbuffer.replace(b'*attr6*', struct.pack('<f', y_scale))
+        flexbuffer = flexbuffer.replace(b'*attr7*', struct.pack('<f', h_scale))
+        flexbuffer = flexbuffer.replace(b'*attr8*', struct.pack('<f', w_scale))
+        flexbuffer = flexbuffer.replace(b'*attr9*', struct.pack('<i', max_detections))
+        flexbuffer = flexbuffer.replace(b'*attr10*', struct.pack('<i', max_classes_per_detection))
+        flexbuffer = flexbuffer.replace(b'*attr11*', struct.pack('<b', use_regular_nms))
+
+        custom_options = builder.CreateNumpyVector(np.array(bytearray(flexbuffer)))
+
+        # operator
+        Operator.OperatorStart(builder)
+        Operator.OperatorAddOpcodeIndex(builder, 0)
+        Operator.OperatorAddInputs(builder, inputs)
+        Operator.OperatorAddOutputs(builder, outputs)
+        Operator.OperatorAddCustomOptions(builder, custom_options)
+        Operator.OperatorAddCustomOptionsFormat(builder, CustomOptionsFormat.FLEXBUFFERS)
+        operator = Operator.OperatorEnd(builder)
+
+        # operators
+        SubGraph.SubGraphStartOperatorsVector(builder, 1)
+        builder.PrependUOffsetTRelative(operator)
+        operators = builder.EndVector(1)
+
+        # subgraph
+        SubGraph.SubGraphStart(builder)
+        SubGraph.SubGraphAddTensors(builder, tensors)
+        SubGraph.SubGraphAddInputs(builder, inputs)
+        SubGraph.SubGraphAddOutputs(builder, outputs)
+        SubGraph.SubGraphAddOperators(builder, operators)
+        subgraph = SubGraph.SubGraphEnd(builder)
+
+        # subgraphs
+        Model.ModelStartSubgraphsVector(builder, 1)
+        builder.PrependUOffsetTRelative(subgraph)
+        subgraphs = builder.EndVector(1)
+
+        description = builder.CreateString("Model for tflite testing")
+
+        Buffer.BufferStartDataVector(builder, 0)
+        data = builder.EndVector(0)
+
+        Buffer.BufferStart(builder)
+        Buffer.BufferAddData(builder, data)
+        buffer = Buffer.BufferEnd(builder)
+
+        Model.ModelStartBuffersVector(builder, 1)
+        builder.PrependUOffsetTRelative(buffer)
+        buffers = builder.EndVector(1)
+
+        # model
+        Model.ModelStart(builder)
+        Model.ModelAddVersion(builder, 3)
+        Model.ModelAddOperatorCodes(builder, op_codes)
+        Model.ModelAddSubgraphs(builder, subgraphs)
+        Model.ModelAddDescription(builder, description)
+        Model.ModelAddBuffers(builder, buffers)
+        model = Model.ModelEnd(builder)
+
+        builder.Finish(model, b"TFL3")
+        return builder.Output()
+
+    def run_tflite_test(self, tflite_model, feed_dict, rtol=1e-07, atol=1e-5):
+        tflite_path = os.path.join(self.test_data_directory, self._testMethodName + ".tflite")
+        dir_name = os.path.dirname(tflite_path)
+        if dir_name:
+            os.makedirs(dir_name, exist_ok=True)
+        with open(tflite_path, 'wb') as f:
+            f.write(tflite_model)
+        tf_lite_output_data, output_names = self.run_tflite(tflite_path, feed_dict)
+
+        g = process_tf_graph(None, opset=self.config.opset,
+                             input_names=list(feed_dict.keys()),
+                             output_names=output_names,
+                             target=self.config.target,
+                             tflite_path=tflite_path)
+        g = optimizer.optimize_graph(g)
+        onnx_from_tfl_output = self.run_backend(g, output_names, feed_dict, postfix="_from_tflite")
+        self.assert_results_equal(tf_lite_output_data, onnx_from_tfl_output, rtol, atol)
diff --git a/tf2onnx/tflite_handlers/__init__.py b/tf2onnx/tflite_handlers/__init__.py
@@ -1,11 +1,12 @@
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT license.
+# SPDX-License-Identifier: Apache-2.0
+
 """tf2onnx.tflite_handlers module"""
 
 from . import (
     tfl_math,
     tfl_nn,
     tfl_controlflow,
     tfl_direct,
-    tfl_tensor
+    tfl_tensor,
+    tfl_postprocess,
 )
diff --git a/tf2onnx/tflite_handlers/tfl_controlflow.py b/tf2onnx/tflite_handlers/tfl_controlflow.py
@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT license.
+# SPDX-License-Identifier: Apache-2.0
+
 
 """
 tfl_controlflow
diff --git a/tf2onnx/tflite_handlers/tfl_direct.py b/tf2onnx/tflite_handlers/tfl_direct.py
@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT license.
+# SPDX-License-Identifier: Apache-2.0
+
 
 """
 tfl_direct
diff --git a/tf2onnx/tflite_handlers/tfl_math.py b/tf2onnx/tflite_handlers/tfl_math.py
@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT license.
+# SPDX-License-Identifier: Apache-2.0
+
 
 """
 tfl_math
diff --git a/tf2onnx/tflite_handlers/tfl_nn.py b/tf2onnx/tflite_handlers/tfl_nn.py
@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT license.
+# SPDX-License-Identifier: Apache-2.0
+
 
 """
 tfl_nn
diff --git a/tf2onnx/tflite_handlers/tfl_postprocess.py b/tf2onnx/tflite_handlers/tfl_postprocess.py
@@ -7,6 +7,7 @@
 
 import logging
 import numpy as np
+from onnx.onnx_pb import TensorProto
 
 from tf2onnx.handler import tfl_op
 from tf2onnx import utils
@@ -76,7 +77,8 @@ def version_11(cls, ctx, node, **kwargs):
 
         nms_inputs = [adjusted_boxes, scores, max_boxes_per_class_const, iou_threshold_const, score_threshold_const]
         # shape: [-1, 3], elts of format [batch_index, class_index, box_index]
-        selected_indices = ctx.make_node('NonMaxSuppression', nms_inputs, attr={'center_point_box': 0}).output[0]
+        selected_indices = ctx.make_node('NonMaxSuppression', nms_inputs, attr={'center_point_box': 0},
+                                         op_name_scope=node.name).output[0]
 
         selected_boxes_idx = GraphBuilder(ctx).make_slice(
             {'data': selected_indices, 'starts': [2], 'ends': [3], 'axes': [1]})
@@ -89,15 +91,23 @@ def version_11(cls, ctx, node, **kwargs):
         box_and_class_idx = ctx.make_node('Concat', [selected_boxes_idx, selected_classes], attr={'axis': 1}).output[0]
 
         box_cnt = ctx.make_node('Shape', [selected_classes_sq]).output[0]
-        box_cnt_float = ctx.make_node('Cast', [box_cnt], attr={'to': box_cnt_dtype}).output[0]
 
         adjusted_boxes_sq = GraphBuilder(ctx).make_squeeze({'data': adjusted_boxes, 'axes': [0]})
         detection_boxes = ctx.make_node('Gather', [adjusted_boxes_sq, selected_boxes_idx_sq]).output[0]
         class_predictions_sq = GraphBuilder(ctx).make_squeeze({'data': class_predictions, 'axes': [0]})
         detection_scores = ctx.make_node('GatherND', [class_predictions_sq, box_and_class_idx]).output[0]
 
         k_const = ctx.make_const(utils.make_name('const_k'), np.array([max_detections], np.int64)).output[0]
-        min_k = ctx.make_node('Min', [k_const, box_cnt]).output[0]
+        if ctx.opset >= 12:
+            min_k = ctx.make_node('Min', [k_const, box_cnt]).output[0]
+        else:
+            # Lower opsets only support Min between floats
+            box_cnt_float = ctx.make_node('Cast', [box_cnt], attr={'to': TensorProto.FLOAT}).output[0]
+            k_const_float = ctx.make_node('Cast', [k_const], attr={'to': TensorProto.FLOAT}).output[0]
+            min_k_float = ctx.make_node('Min', [k_const_float, box_cnt_float]).output[0]
+            min_k = ctx.make_node('Cast', [min_k_float], attr={'to': TensorProto.INT64}).output[0]
+        min_k_cast = ctx.make_node('Cast', [min_k], attr={'to': box_cnt_dtype}).output[0]
+
         scores_top_k, scores_top_k_idx = ctx.make_node('TopK', [detection_scores, min_k], output_count=2).output
 
         scores_top_k_idx_unsq = GraphBuilder(ctx).make_unsqueeze({'data': scores_top_k_idx, 'axes': [0]})
@@ -107,7 +117,7 @@ def version_11(cls, ctx, node, **kwargs):
         classes_sort_cast = ctx.make_node('Cast', [selected_classes_sort], attr={'to': classes_dtype}).output[0]
         detection_boxes_sorted = ctx.make_node('Gather', [detection_boxes, scores_top_k_idx_unsq]).output[0]
 
-        pad_amount = ctx.make_node('Sub', [k_const, box_cnt]).output[0]
+        pad_amount = ctx.make_node('Sub', [k_const, min_k]).output[0]
 
         quad_zero_const = ctx.make_const(utils.make_name('quad_zero_const'), np.array([0, 0, 0, 0], np.int64)).output[0]
         duo_zero_const = ctx.make_const(utils.make_name('duo_zero_const'), np.array([0, 0], np.int64)).output[0]
@@ -123,4 +133,6 @@ def version_11(cls, ctx, node, **kwargs):
         ctx.replace_all_inputs(node.output[0], detection_boxes_padded)
         ctx.replace_all_inputs(node.output[1], detection_classes_padded)
         ctx.replace_all_inputs(node.output[2], detection_scores_padded)
-        ctx.replace_all_inputs(node.output[3], box_cnt_float)
+        ctx.replace_all_inputs(node.output[3], min_k_cast)
+
+        ctx.remove_node(node.name)
diff --git a/tf2onnx/tflite_handlers/tfl_tensor.py b/tf2onnx/tflite_handlers/tfl_tensor.py
@@ -1,5 +1,5 @@
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT license.
+# SPDX-License-Identifier: Apache-2.0
+
 
 """
 tfl_tensor
