Made TFLite layer creation explicit

Change-Id: I7fbf6a5a2163c1fada49477f86d84f1bc09bd57c

Author: Ashutosh Parkhi

python/tvm/relay/testing/tflite.py

@@ -35,44 +35,53 @@ def __init__(self, dtype):
         self.shape_dict = {}
         self.dtype_dict = {}
 
-    @tf.function
-    def conv2d_single_function(self, ifm_tensor, args):
-        """Returns TFLite Conv2d layer"""
-        assert len(args) == 6, "Conv2D needs (ifm_shape, kernel_shape, strides, padding, dilation)"
-        _, kernel_shape, strides, padding, dilation, activation = args
-        op = tf.nn.conv2d(
-            ifm_tensor,
-            filters=tf.constant(
-                np.random.uniform(size=[kernel_shape[0], kernel_shape[1], 3, 3]),
-                dtype=tf.float32,
-            ),
-            strides=[1, strides[0], strides[1], 1],
-            padding=padding,
-            dilations=dilation,
-        )
-        if activation == "RELU":
-            op = tf.nn.relu(op)
-        elif activation == "NONE":
-            pass
-        else:
-            assert False, "Unsupported activation {}".format(activation)
-        return op
-
-    def create_tflite_model(self, op_type, *args):
-        """Returns TFLite serial graph, Relay module, Relay params based on op_type"""
-        concrete_func = None
-        input_shape = None
-        if op_type == "conv2d_single":
-            input_shape = args[0]
-            ifm_tensor = tf.TensorSpec(input_shape, dtype=tf.float32, name="input")
-            concrete_func = self.conv2d_single_function.get_concrete_function(ifm_tensor, args)
-        else:
-            assert False, "Unsupported op_type {}".format(op_type)
+    def create_conv2d_single(self, kernel_shape, strides, padding, dilation, activation):
+        @tf.function
+        def conv2d_single_function(ifm_tensor):
+            """Returns TFLite Conv2d layer"""
+            op = tf.nn.conv2d(
+                ifm_tensor,
+                filters=tf.constant(
+                    np.random.uniform(size=[kernel_shape[0], kernel_shape[1], 3, 3]),
+                    dtype=tf.float32,
+                ),
+                strides=[1, strides[0], strides[1], 1],
+                padding=padding,
+                dilations=dilation,
+            )
+            if activation == "RELU":
+                op = tf.nn.relu(op)
+            elif activation == "NONE":
+                pass
+            else:
+                assert False, "Unsupported activation {}".format(activation)
+            return op
+
+        return conv2d_single_function
+
+    def create_tflite_model(self, tfl_function, shapes, ranges=None):
+        """Creates TFLite serial graph"""
+        tensor_specs = []
+        for i, shape in enumerate(shapes):
+            input_name = "input_" + str(i)
+            self.shape_dict.update({input_name: shape})
+            self.dtype_dict.update({input_name: self.dtype})
+            tensor_specs.append(tf.TensorSpec(shape, dtype=tf.float32, name=input_name))
+        concrete_func = tfl_function.get_concrete_function(*tensor_specs)
+
+        if not ranges:
+            ranges = [(0, 1) for _ in shapes]
 
         def representative_dataset():
             for _ in range(100):
-                data = np.random.rand(*tuple(input_shape))
-                yield [data.astype(np.float32)]
+                inputs = []
+                for i, shape in enumerate(shapes):
+                    data = np.random.uniform(
+                        low=ranges[i][0], high=ranges[i][1], size=tuple(shape)
+                    ).astype("float32")
+                    inputs.append(data)
+
+                yield inputs
 
         converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
         converter.optimizations = [tf.lite.Optimize.DEFAULT]
@@ -81,8 +90,6 @@ def representative_dataset():
         converter.inference_input_type = tf.int8
         converter.inference_output_type = tf.int8
         self.serial_model = converter.convert()
-        self.shape_dict = {"input": input_shape}
-        self.dtype_dict = {"input": self.dtype}
 
     def convert_to_relay(self):
         """Converts TFLite serialized graph into Relay"""

tests/python/contrib/test_cmsisnn/test_conv2d.py

@@ -317,9 +317,10 @@ def test_conv2d_int8_tflite(ifm_shape, kernel_shape, strides, dilation, padding,
     from tvm.relay.testing.tflite import TFLiteModel
 
     tfl_model = TFLiteModel(dtype)
-    tfl_model.create_tflite_model(
-        "conv2d_single", ifm_shape, kernel_shape, strides, padding, dilation, activation
+    conv2d_function = tfl_model.create_conv2d_single(
+        kernel_shape, strides, padding, dilation, activation
     )
+    tfl_model.create_tflite_model(conv2d_function, [ifm_shape])
     relay_mod, relay_params = tfl_model.convert_to_relay()
 
     cmsisnn_mod = cmsisnn.partition_for_cmsisnn(relay_mod, relay_params)

tests/python/contrib/test_cmsisnn/utils.py

@@ -225,134 +225,3 @@ def make_qnn_relu(expr, fused_activation_fn, scale, zero_point, dtype):
         )
     if fused_activation_fn == "RELU":
         return tvm.relay.op.clip(expr, a_min=max(qmin, quantize(0.0)), a_max=qmax)
-
-
-def generate_random_input_data(seed, shape, dtype):
-    """
-    Generates randomized input numpy arrays based on shape and dtype
-    """
-    random_state = np.random.RandomState(seed)
-    if dtype == np.float32:
-        return random_state.uniform(-1, 1, size).astype(dtype)
-    else:
-        low = np.iinfo(dtype).min
-        high = np.iinfo(dtype).max + 1
-        return random_state.randint(low, high, shape, dtype)
-
-
-def generate_ref_data_tflite(model):
-    """
-    This method uses TFLite reference kernels to generate reference output.
-    Random input generator is used to get the input data.
-    It returns randomized inputs and reference outputs.
-    """
-    import tensorflow as tf
-    from distutils.version import LooseVersion
-
-    output_tolerance = None
-    if tf.__version__ < LooseVersion("2.5.0"):
-        output_tolerance = 1
-        interpreter = tf.lite.Interpreter(model_content=model)
-    else:
-        from tensorflow.lite.python.interpreter import OpResolverType
-
-        output_tolerance = 0
-        interpreter = tf.lite.Interpreter(
-            model_content=model,
-            experimental_op_resolver_type=OpResolverType.BUILTIN_REF,
-            experimental_preserve_all_tensors=False,
-        )
-
-    interpreter.allocate_tensors()
-    input_details = interpreter.get_input_details()
-    output_details = interpreter.get_output_details()
-
-    # Generate predictable randomized input
-    seed = 0
-    input_data = {}
-    for input_detail in input_details:
-        input_values = generate_random_input_data(
-            seed, input_detail["shape"], input_detail["dtype"]
-        )
-        interpreter.set_tensor(input_detail["index"], input_values)
-        input_data.update({input_detail["name"]: input_values})
-
-    interpreter.invoke()
-
-    # Obtain the expected output from interpreter
-    expected_output_data = {}
-    for output_detail in output_details:
-        expected_output_data.update(
-            {output_detail["name"]: interpreter.get_tensor(output_detail["index"])}
-        )
-
-    return input_data, expected_output_data, output_tolerance
-
-
-def create_conv2d_tflite_model(ifm_shape, kernel_shape, strides, dilation, padding, activation):
-    """This method prepares TFlite graph with a single Conv2d layer"""
-    import tensorflow as tf
-
-    class Model(tf.Module):
-        @tf.function
-        def tf_function(self, x):
-            # Use tf.nn API to create the model
-            tf_strides = [1, strides[0], strides[1], 1]
-            op = tf.nn.conv2d(
-                x,
-                filters=tf.constant(
-                    np.random.uniform(size=[kernel_shape[0], kernel_shape[1], 3, 3]),
-                    dtype=tf.float32,
-                ),
-                strides=tf_strides,
-                padding=padding,
-                dilations=dilation,
-            )
-            if activation:
-                op = tf.nn.relu(op)
-            return op
-
-    model = Model()
-    concrete_func = model.tf_function.get_concrete_function(
-        tf.TensorSpec(ifm_shape, dtype=tf.float32)
-    )
-
-    def representative_dataset():
-        for _ in range(100):
-            data = np.random.rand(*tuple(ifm_shape))
-            yield [data.astype(np.float32)]
-
-    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
-    converter.optimizations = [tf.lite.Optimize.DEFAULT]
-    converter.representative_dataset = representative_dataset
-    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
-    converter.inference_input_type = tf.int8
-    converter.inference_output_type = tf.int8
-    tflite_model = converter.convert()
-    return tflite_model
-
-
-def create_conv2d_tflite_relay_models(
-    ifm_shape, kernel_shape, strides, dilation, padding, activation, dtype
-):
-    """
-    This method creates a conv2d TFLite layer and prepared TFLite model from it.
-    Converts that into the Relay module and params.
-    Returns TFLite model, Relay module and params.
-    """
-    pytest.importorskip("tflite")
-    import tflite.Model
-
-    serialized_tflite_model = create_conv2d_tflite_model(
-        ifm_shape, kernel_shape, strides, dilation, padding, activation
-    )
-
-    tflite_model = tflite.Model.Model.GetRootAsModel(serialized_tflite_model, 0)
-
-    relay_module, params = relay.frontend.from_tflite(
-        tflite_model,
-        shape_dict={"input": ifm_shape},
-        dtype_dict={"input": dtype},
-    )
-
-    return serialized_tflite_model, relay_module, params