From d61be797d8770596f5436aa32561d0e411b18616 Mon Sep 17 00:00:00 2001 From: Shai Maor Date: Tue, 26 Oct 2021 16:13:28 +0300 Subject: [PATCH] Support quantised Squared Difference operator in TFLite --- python/tvm/relay/frontend/tflite.py | 20 ++++++++++++-------- tests/python/frontend/tflite/test_forward.py | 8 ++++++-- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index b0b2bd3e75d9..45dd5a184fa1 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -1243,7 +1243,7 @@ def convert_square(self, op): return out - def _convert_elemwise(self, relay_op, op, ignore_qnn_params=False): + def _convert_elemwise(self, relay_op, op, ignore_qnn_params=False, dequantize=False): """Generic method to Convert TFLite elemwise""" try: from tflite.AddOptions import AddOptions @@ -1277,8 +1277,13 @@ def _convert_elemwise(self, relay_op, op, ignore_qnn_params=False): # If quantized, extracts qnn params and call QNN add operator. if not ignore_qnn_params and lhs_tensor.qnn_params: - assert rhs_tensor.qnn_params, "Both tensors should be quantized." - assert output_tensor.qnn_params, "Output tensor should be quantized." + if not dequantize: + assert rhs_tensor.qnn_params, "Both tensors should be quantized." + assert output_tensor.qnn_params, "Output tensor should be quantized." + else: + lhs_expr = self.dequantize(lhs_expr, lhs_tensor) + rhs_expr = self.dequantize(rhs_expr, rhs_tensor) + out = relay_op( lhs=lhs_expr, rhs=rhs_expr, @@ -1292,6 +1297,9 @@ def _convert_elemwise(self, relay_op, op, ignore_qnn_params=False): else: out = relay_op(lhs_expr, rhs_expr) + if dequantize and output_tensor.qnn_params: + out = self.quantize(out, output_tensor) + # Options (fused_activation_function) options = None if op.BuiltinOptionsType() == BuiltinOptions.AddOptions: @@ -1393,11 +1401,7 @@ def convert_greater(self, op): def convert_squared_difference(self, op): """Convert TFLite SQUARED DIFFERENCE""" # Check if the input tensor is quantized, call QNN op - if self.is_quantized(op): - raise tvm.error.OpNotImplemented( - "TFlite quantized squared difference operator is not supported yet." - ) - difference = self._convert_elemwise(_op.subtract, op) + difference = self._convert_elemwise(_op.subtract, op, dequantize=True) # _convert_elemwise has guaranteed only have one output tensor exp_type = self.get_tensor_type_str(self.get_output_tensors(op)[0].tensor.Type()) out = _op.power(difference, relay.const(2, exp_type)) diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py index d234cd13c986..b92bbdef32a6 100644 --- a/tests/python/frontend/tflite/test_forward.py +++ b/tests/python/frontend/tflite/test_forward.py @@ -2227,9 +2227,11 @@ def _test_not_equal(data): # ------------------ -def _test_squared_difference(data): +def _test_squared_difference(data, fused_activation_function=None, quantized=False, qnn_op=None): """One iteration of squared difference""" - return _test_elemwise(math_ops.squared_difference, data) + return _test_elemwise( + math_ops.squared_difference, data, fused_activation_function, quantized, qnn_op + ) ####################################################################### @@ -2293,6 +2295,7 @@ def _test_elemwise_qnn_out_range(qnn_op): _test_mul: (-5e3, 5e3), _test_maximum: (-112, 111), _test_minimum: (-128, 127), + _test_squared_difference: (0, 225e2), } return qnn_out_range[qnn_op] @@ -2323,6 +2326,7 @@ def test_all_elemwise(): _test_forward_elemwise_quantized(_test_minimum) _test_forward_elemwise(_test_greater) _test_forward_elemwise(_test_squared_difference) + _test_forward_elemwise_quantized(_test_squared_difference) _test_forward_elemwise(_test_greater_equal) _test_forward_elemwise(_test_less) _test_forward_elemwise(_test_less_equal)