tvm-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From GitBox <...@apache.org>
Subject [GitHub] [incubator-tvm] apivovarov commented on a change in pull request #4447: [Relay][Frontend][TFlite] Add parses support for UNPACK tflite operator
Date Wed, 04 Dec 2019 23:27:41 GMT
apivovarov commented on a change in pull request #4447: [Relay][Frontend][TFlite] Add parses
support for UNPACK tflite operator
URL: https://github.com/apache/incubator-tvm/pull/4447#discussion_r354039512
 
 

 ##########
 File path: python/tvm/relay/frontend/tflite.py
 ##########
 @@ -1238,6 +1239,54 @@ def convert_pack(self, op):
         out = _op.concatenate(in_exprs_reshaped, pack_axis)
         return out
 
+    def convert_unpack(self, op):
+        """Convert TFLite unpack"""
+        try:
+            from tflite.BuiltinOptions import BuiltinOptions
+            from tflite.Operator import Operator
+            from tflite.UnpackOptions import UnpackOptions
+        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) == 1, "input tensors length should be 1"
+        input_tensor = input_tensors[0]
+        in_expr = self.get_expr(input_tensor.tensor_idx)
+        assert op.BuiltinOptionsType() == BuiltinOptions.UnpackOptions
+        op_options = op.BuiltinOptions()
+        unpack_options = UnpackOptions()
+        unpack_options.Init(op_options.Bytes, op_options.Pos)
+        num_unpacks = unpack_options.Num()
+        unpack_axis = unpack_options.Axis()
+
+        # Relay doesn't support 'unpack' operator so we use 'split' & 'squeeze' instead.
+        # We have to do 'squeeze' along the split axis but Relay expects
+        # squeeze_axis to be either None or List
+        if unpack_axis == 0:
+            squeeze_axis = None
+        else:
+            squeeze_axis = [unpack_axis]
+
+        # Relay doesn't like TupleWrapper of 1 element so we isolate the case of unpacking
+        # a tensor by an axis with len(axis) == 1. For reference see convert_split()
+        # Such unpacking will result in the same tensor so we omit 'split' and only squeeze
+        # along the axis of dim == 1
+        if num_unpacks == 1:
+            squeezed = _op.squeeze(in_expr, axis=squeeze_axis)
+            if isinstance(squeezed, _expr.TupleWrapper):
+                squeezed = squeezed[0]
+        else:
+            splitted = _op.split(in_expr,
+                                 indices_or_sections=num_unpacks,
+                                 axis=unpack_axis)
+            squeezed = _expr.TupleWrapper(
+                _expr.Tuple([_op.squeeze(split_item, axis=squeeze_axis) \
+                             for split_item in splitted]), len(splitted))
+        out = squeezed
+
+        return out
 
 Review comment:
   just "return squeezed"?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

Mime
View raw message