Add Reduce operators to TFLite (#3421)

diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
index acfa914..fe16387 100644 (file)
--- a/python/tvm/relay/frontend/tflite.py
+++ b/python/tvm/relay/frontend/tflite.py
@@ -73,6 +73,10 @@ class OperatorConverter(object):
             'POW': self.convert_pow,
             'MAXIMUM': self.convert_maximum,
             'MINIMUM': self.convert_minimum,
+            'REDUCE_MIN': self._convert_reduce_min,
+            'REDUCE_MAX': self._convert_reduce_max,
+            'MEAN': self._convert_reduce_mean,
+            'REDUCE_PROD': self._convert_reduce_prod,
             'FULLY_CONNECTED': self.convert_fully_connected,
             'PAD': self.convert_pad,
             'LOGISTIC': self.convert_logistic,
@@ -427,6 +431,48 @@ class OperatorConverter(object):
     def convert_minimum(self, op):
         return self._convert_elemwise(_op.minimum, op)
 
+    def _convert_reduce(self, relay_op, op):
+        """Generic method to Convert TFLite MEAN operators"""
+        try:
+            from tflite.BuiltinOptions import BuiltinOptions
+            from tflite.Operator import Operator
+            from tflite.ReducerOptions import ReducerOptions
+        except ImportError:
+            raise ImportError("The tflite package must be installed")
+
+        assert isinstance(op, Operator)
+        input_tensors = self.get_input_tensors(op)
+        assert len(input_tensors) == 2, "input tensors length should be 2"
+
+        # input_tensor
+        input_tensor = input_tensors[0]
+        in_expr = self.get_expr(input_tensor.tensor_idx)
+
+        # axis
+        axis = tuple(self.get_tensor_value(input_tensors[1]))
+
+        # Options - keep_dims (bool)
+        assert op.BuiltinOptionsType() == BuiltinOptions.ReducerOptions
+        reduce_options = ReducerOptions()
+        op_options = op.BuiltinOptions()
+        reduce_options.Init(op_options.Bytes, op_options.Pos)
+        keep_dims = reduce_options.KeepDims()
+
+        out = relay_op(in_expr, axis, keep_dims)
+        return out
+
+    def _convert_reduce_min(self, op):
+        return self._convert_reduce(_op.reduce.min, op)
+
+    def _convert_reduce_max(self, op):
+        return self._convert_reduce(_op.reduce.max, op)
+
+    def _convert_reduce_mean(self, op):
+        return self._convert_reduce(_op.reduce.mean, op)
+
+    def _convert_reduce_prod(self, op):
+        return self._convert_reduce(_op.reduce.prod, op)
+
     def convert_fully_connected(self, op):
         """Convert TFLite fully connected"""
         try:

diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
index 0a011a6..41147fe 100644 (file)
--- a/tests/python/frontend/tflite/test_forward.py
+++ b/tests/python/frontend/tflite/test_forward.py
@@ -360,7 +360,7 @@ def test_forward_concatenation():
 # ---
 
 def _test_elemwise(math_op, data, fused_activation_function=None):
-    """ One iteration of add """
+    """ One iteration of elemwise """
 
     assert len(data) == 2
 
@@ -458,6 +458,74 @@ def test_all_elemwise():
     _test_forward_elemwise(_test_minimum)
 
 #######################################################################
+# Reduce
+# ------
+
+def _test_reduce(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=data[0].dtype, name='in')
+        out = math_op(in_data, data[1], keep_dims)
+        compare_tflite_with_tvm([data[0]], ['in:0'], [in_data], [out])
+
+
+#######################################################################
+# Reduce_min
+# ----------
+
+def _test_reduce_min(data, keep_dims=None):
+    """ One iteration of reduce_min """
+    return _test_reduce(math_ops.reduce_min, data, keep_dims)
+
+#######################################################################
+# Reduce_max
+# ----------
+
+def _test_reduce_max(data, keep_dims=None):
+    """ One iteration of reduce_max """
+    return _test_reduce(math_ops.reduce_max, data, keep_dims)
+
+#######################################################################
+# Reduce_mean
+# -----------
+
+def _test_reduce_mean(data, keep_dims=None):
+    """ One iteration of reduce_mean """
+    return _test_reduce(math_ops.reduce_mean, data, keep_dims)
+
+#######################################################################
+# Reduce_prod
+# -----------
+
+def _test_reduce_prod(data, keep_dims=None):
+    """ One iteration of reduce_prod """
+    return _test_reduce(math_ops.reduce_prod, data, keep_dims)
+
+
+def _test_forward_reduce(testop):
+    """ Reduce """
+    data0 = [np.random.rand(16, 16, 16, 16).astype("float32"), None]
+    data1 = [np.random.rand(16, 16, 16, 16).astype("float32"), np.array([1, 2], dtype=np.int32)]
+    testop(data0)
+    testop(data0, keep_dims=False)
+    testop(data0, keep_dims=True)
+    testop(data1)
+    testop(data1, keep_dims=False)
+    testop(data1, keep_dims=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(_test_reduce_prod)
+
+
+#######################################################################
 # Squeeze
 # -------
 
@@ -695,6 +763,9 @@ if __name__ == '__main__':
     # Elemwise
     test_all_elemwise()
 
+    # Reduce
+    test_all_reduce()
+
     # End to End
     test_forward_mobilenet_v1()
     test_forward_mobilenet_v2()