Added tflite frontend support for quantized mean. (#4339)

diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
index f00a6a0..e2dc0e7 100644 (file)
--- a/python/tvm/relay/frontend/tflite.py
+++ b/python/tvm/relay/frontend/tflite.py
@@ -658,7 +658,24 @@ class OperatorConverter(object):
         reduce_options.Init(op_options.Bytes, op_options.Pos)
         keep_dims = reduce_options.KeepDims()
 
+        if input_tensor.qnn_params:
+            in_expr = _op.cast(in_expr, "int32")
+
         out = relay_op(in_expr, axis, keep_dims)
+
+        # Finally if the reduce is quantized. Add a requantize at the end.
+        output_tensors = self.get_output_tensors(op)
+        assert len(output_tensors) == 1, "output tensors length should be 1"
+        output_tensor = output_tensors[0]
+        output_tensor_type_str = self.get_tensor_type_str(output_tensor.tensor.Type())
+        if output_tensor.qnn_params:
+            out = _qnn.op.requantize(out,
+                                     input_scale=input_tensor.qnn_params['scale'],
+                                     input_zero_point=input_tensor.qnn_params['zero_point'],
+                                     output_scale=output_tensor.qnn_params['scale'],
+                                     output_zero_point=output_tensor.qnn_params['zero_point'],
+                                     out_dtype=output_tensor_type_str)
+
         return out
 
     def _convert_reduce_min(self, op):

diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
index 8292dd5..555dc57 100644 (file)
--- a/tests/python/frontend/tflite/test_forward.py
+++ b/tests/python/frontend/tflite/test_forward.py
@@ -722,6 +722,19 @@ def _test_reduce(math_op, data, keep_dims=None):
         out = math_op(in_data, data[1], keep_dims)
         compare_tflite_with_tvm([data[0]], ['in:0'], [in_data], [out])
 
+def _test_reduce_quantize(math_op, data, keep_dims=None):
+    """ One iteration of reduce """
+
+    assert len(data) == 2
+
+    # Test with tensor and constant
+    with tf.Graph().as_default():
+        in_data = [array_ops.placeholder(shape=data[0].shape, dtype="float32", name='in')]
+        inq_data = [tf.quantization.fake_quant_with_min_max_args(in_data[0], min=-100, max=100, name="inq_0")]
+        out = math_op(inq_data, data[1], keep_dims)
+        out = tf.quantization.fake_quant_with_min_max_args(out, min=-200, max=200, name="out")
+        compare_tflite_with_tvm([data[0]], ['inq_0:0'], [inq_data[0]], [out], quantized=True)
+
 
 #######################################################################
 # Reduce_min
@@ -743,9 +756,12 @@ def _test_reduce_max(data, keep_dims=None):
 # Reduce_mean
 # -----------
 
-def _test_reduce_mean(data, keep_dims=None):
+def _test_reduce_mean(data, keep_dims=None, quantized=False):
     """ One iteration of reduce_mean """
-    return _test_reduce(math_ops.reduce_mean, data, keep_dims)
+    if quantized:
+        return _test_reduce_quantize(math_ops.reduce_mean, data, keep_dims)
+    else:
+        return _test_reduce(math_ops.reduce_mean, data, keep_dims)
 
 #######################################################################
 # Reduce_prod
@@ -775,11 +791,17 @@ def _test_forward_reduce(testop):
     testop(data1, keep_dims=False)
     testop(data1, keep_dims=True)
 
+def _test_forward_reduce_quantized(testop):
+    data0 = [np.array(np.random.uniform(0, 255, (3, 6)), dtype=np.uint8), np.array([1, 2], dtype=np.int32)]
+    testop(data0, quantized=True)
+    testop(data0, keep_dims=False, quantized=True)
+    testop(data0, keep_dims=True, quantized=True)
 
 def test_all_reduce():
     _test_forward_reduce(_test_reduce_min)
     _test_forward_reduce(_test_reduce_max)
     _test_forward_reduce(_test_reduce_mean)
+    _test_forward_reduce_quantized(_test_reduce_mean)
     _test_forward_reduce(_test_reduce_prod)
     _test_forward_reduce(_test_reduce_sum)