From 0061bae7354c33053dfcdf2a9082b81094d1bccc Mon Sep 17 00:00:00 2001 From: Alexander Pivovarov Date: Sat, 22 Jun 2019 21:18:46 -0700 Subject: [PATCH] Add Reduce operators to TFLite --- python/tvm/relay/frontend/tflite.py | 46 ++++++++++++ tests/python/frontend/tflite/test_forward.py | 73 +++++++++++++++++++- 2 files changed, 118 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index acfa9148a756..fe163871fa60 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -73,6 +73,10 @@ def __init__(self, model, subgraph, exp_tab): '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 @@ def convert_maximum(self, op): 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 0a011a6e709b..41147fe9e9bf 100644 --- 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 @@ -457,6 +457,74 @@ def test_all_elemwise(): _test_forward_elemwise(_test_maximum) _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 @@ def test_forward_ssd_mobilenet_v1(): # Elemwise test_all_elemwise() + # Reduce + test_all_reduce() + # End to End test_forward_mobilenet_v1() test_forward_mobilenet_v2()