diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index f00a6a02ab6c..e2dc0e77d980 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -658,7 +658,24 @@ def _convert_reduce(self, relay_op, op): 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 8292dd5bdae6..555dc579b0d8 100644 --- 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)