Adding support for TFLite QnnSub operator. (#5230)

apache · Apr 10, 2020 · 6ecfaaf · 6ecfaaf
1 parent 029388f
commit 6ecfaaf
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
@@ -926,8 +926,7 @@ def convert_sub(self, op):
         """Convert TFLite SUB"""
         # Check if the input tensor is quantized, call QNN op
         if self.is_quantized(op):
-            raise tvm.error.OpNotImplemented(
-                'TFlite quantized SUB operator is not supported yet.')
+            return self._convert_elemwise(_qnn.op.subtract, op)
         return self._convert_elemwise(_op.subtract, op)
 
     def convert_mul(self, op):
@@ -1355,7 +1354,9 @@ def convert_conv(self, op, conv_type):
         if is_depthwise_conv:
             params['channels'] = int(in_channels)
             params['groups'] = int(input_c)
-            params['kernel_layout'] = 'HWOI'
+            # If number of input channels is 1, treat as normal
+            # convolution.
+            params['kernel_layout'] = 'HWIO' if input_c == 1 else 'HWOI'
         else:
             params['channels'] = int(output_channels)
             params['kernel_layout'] = 'HWIO'

diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
@@ -541,6 +541,8 @@ def test_forward_convolution():
     _test_convolution([4, 17, 17, 124], [1, 1, 124, 1], [1, 1], [1, 1], 'SAME', 'NHWC', True)
     _test_convolution([4, 17, 17, 12], [3, 3, 12, 1], [1, 1], [2, 2], 'VALID', 'NHWC', True)
     _test_convolution([4, 17, 17, 12], [3, 3, 12, 2], [1, 1], [2, 2], 'VALID', 'NHWC', True)
+    # dephtwise convolution with single input channel
+    _test_convolution([1, 76, 64, 1], [9, 5, 1, 96], [1, 1], [1, 1], 'SAME', 'NHWC', True)
 
 
 #######################################################################
@@ -902,9 +904,9 @@ def _test_add(data, fused_activation_function=None, quantized=False, qnn_op=None
 # Subtract
 # --------
 
-def _test_sub(data, fused_activation_function=None):
+def _test_sub(data, fused_activation_function=None, quantized=False, qnn_op=None):
     """ One iteration of subtract """
-    return _test_elemwise(math_ops.subtract, data, fused_activation_function)
+    return _test_elemwise(math_ops.subtract, data, fused_activation_function, quantized, qnn_op)
 #######################################################################
 # Mul
 # ---
@@ -1036,6 +1038,7 @@ def test_all_elemwise():
     _test_forward_elemwise(partial(_test_add, fused_activation_function="RELU"))
     _test_forward_elemwise(partial(_test_add, fused_activation_function="RELU6"))
     _test_forward_elemwise(_test_sub)
+    _test_forward_elemwise_quantized(_test_sub)
     _test_forward_elemwise(partial(_test_sub, fused_activation_function="RELU"))
     _test_forward_elemwise(partial(_test_sub, fused_activation_function="RELU6"))
     _test_forward_elemwise(_test_mul)