OP: enable Resize/Prelu/Sigmoid, modified Reshape

tests/test_ops.py

@@ -38,6 +38,7 @@ def test_ops_implicit_layout():
         'fullyconnected.float32',
         'fullyconnected-relu6.float32',
         'maxpooling.float32',
+        'resize-bilinear.float32',
     )
 
     for op in OP_LIST:
@@ -73,6 +74,8 @@ def test_ops_layout_transparent():
         'mul.float32',
         'relu6.float32',
         'relu.float32',
+        'prelu.float32',
+        'sigmoid.float32',
     )
 
     for op in OP_LIST:

tflite2onnx/op/__init__.py

@@ -1,4 +1,5 @@
 from tflite2onnx.op.activation import ReLU
+from tflite2onnx.op.activation import Logistic
 from tflite2onnx.op.binary import Binary
 from tflite2onnx.op.common import OpFactory
 from tflite2onnx.op.common import Operator  # noqa: F401
@@ -10,6 +11,7 @@
 from tflite2onnx.op.quantize import Quantize
 from tflite2onnx.op.reduce import Reduce
 from tflite2onnx.op.reshape import Reshape
+from tflite2onnx.op.resize import Resize
 from tflite2onnx.op.slice import Slice
 from tflite2onnx.op.softmax import Softmax
 from tflite2onnx.op.split import Split
@@ -24,8 +26,10 @@
 OpFactory.register(Pooling)
 OpFactory.register(Quantize)
 OpFactory.register(ReLU)
+OpFactory.register(Logistic)
 OpFactory.register(Reduce)
 OpFactory.register(Reshape)
+OpFactory.register(Resize)
 OpFactory.register(Slice)
 OpFactory.register(Softmax)
 OpFactory.register(Split)

tflite2onnx/op/activation.py

@@ -1,15 +1,51 @@
 import logging
 import tflite
 
+from tflite2onnx import mapping
 from tflite2onnx.op.common import Operator
 
 logger = logging.getLogger('tflite2onnx')
 
 
+class Logistic(Operator):
+    TypeMapping = {
+        tflite.BuiltinOperator.LOGISTIC: 'Sigmoid',
+    }
+
+    def __init__(self, TFactory, index):
+        super().__init__(TFactory, index)
+        self.setInited()
+
+    @property
+    def type(self):
+        return 'Sigmoid'
+
+    def parse(self):
+        logger.debug("Parsing %s...", self.type)
+
+        op = self.tflite
+        opcode = self.model.OperatorCodes(op.OpcodeIndex()).BuiltinCode()
+        assert (opcode in self.TypeMapping)
+
+        assert (op.InputsLength() == 1)
+        assert (op.OutputsLength() == 1)
+        self.parseInput(0)
+        self.parseOutput(0)
+
+        self.setParsed()
+
+    def propagatableTensors(self):
+        return self.inputs + self.outputs
+
+    def transform(self):
+        pass
+
+
 class ReLU(Operator):
     TypeMapping = {
         tflite.BuiltinOperator.RELU: 'Relu',
         tflite.BuiltinOperator.RELU6: 'Clip',
+        tflite.BuiltinOperator.PRELU: 'PRelu',
     }
 
     def __init__(self, TFactory, index, preset_opcode=None):
@@ -40,7 +76,10 @@ def parse(self):
         opcode = self.model.OperatorCodes(op.OpcodeIndex()).BuiltinCode()
         assert(opcode in self.TypeMapping)
 
-        assert(op.InputsLength() == 1)
+        if opcode == tflite.BuiltinOperator.PRELU:
+            assert (op.InputsLength() == 2)
+        else:
+            assert(op.InputsLength() == 1)
         assert(op.OutputsLength() == 1)
 
         it = self.parseInput(0)
@@ -53,6 +92,15 @@ def parse(self):
             tmax.addConsumer(self)
             self.inputs.append(tmax)
 
+        if opcode == tflite.BuiltinOperator.PRELU:
+            # `alpha` should be a learned array with the same shape as `X`
+            # But there is no `batch_size` dimension in its shape,
+            # which will cause `out of index` exception during axis transform
+            # so we copy the `batch_size` from `X` to expand its dimension
+            alpha = self.parseInput(1)
+            alpha.shape.insert(0, self.inputs[0].shape[0])
+            alpha.dtype = mapping.DTYPE_NAME2ONNX['float32']
+
         self.parseOutput(0)
 
         self.setParsed()

tflite2onnx/op/reshape.py

@@ -32,24 +32,38 @@ def parse(self):
         opcode = self.model.OperatorCodes(op.OpcodeIndex()).BuiltinCode()
         assert(opcode is tflite.BuiltinOperator.RESHAPE)
 
-        assert(op.InputsLength() == 2)
+        assert(op.InputsLength() >= 1)
         assert(op.OutputsLength() == 1)
 
         # input
-        self.parseInput(0)
-
-        # shape
-        st = self.parseInput(1)
-        # TFLite shape is int32 data type, ONNX is int64
-        st.dtype = mapping.DTYPE_NAME2ONNX['int64']
-        if st.isInitializer:
-            st.data = st.data.astype('int64')
-        if (len(st.shape) > 1):
-            logger.warning("ONNXRuntime doesn't support 2+rank shape, "
-                           "flatten if the shape is initialzier, ignore otherwise."
-                           "https://github.com/jackwish/tflite2onnx/issues/16")
+        data = self.parseInput(0)
+
+        if op.InputsLength() == 1:
+            # options
+            op_opt = op.BuiltinOptions()
+            option = tflite.ReshapeOptions()
+            option.Init(op_opt.Bytes, op_opt.Pos)
+            sp = option.NewShapeAsNumpy()
+            sp = self.TFactory.createVector(data, sp)
+            sp.parse()
+            sp.addConsumer(self)
+            sp.dtype = mapping.DTYPE_NAME2ONNX['int64']
+            self.inputs.append(sp)
+
+        if op.InputsLength() == 2:
+            # shape
+            st = self.parseInput(1)
+
+            # TFLite shape is int32 data type, ONNX is int64
+            st.dtype = mapping.DTYPE_NAME2ONNX['int64']
             if st.isInitializer:
-                st.shape = (np.prod(np.array(st.shape)),)
+                st.data = st.data.astype('int64')
+            if len(st.shape) > 1:
+                logger.warning("ONNXRuntime doesn't support 2+rank shape, "
+                               "flatten if the shape is initialzier, ignore otherwise."
+                               "https://github.com/jackwish/tflite2onnx/issues/16")
+                if st.isInitializer:
+                    st.shape = (np.prod(np.array(st.shape)),)
 
         # output
         self.parseOutput(0)

tflite2onnx/op/resize.py

@@ -0,0 +1,118 @@
+import logging
+import tflite
+import numpy as np
+
+from tflite2onnx import mapping
+from tflite2onnx.op.common import Operator
+
+logger = logging.getLogger('tflite2onnx')
+
+class Resize(Operator):
+    TypeMapping = {
+        tflite.BuiltinOperator.RESIZE_NEAREST_NEIGHBOR: 'Resize',
+        tflite.BuiltinOperator.RESIZE_BILINEAR: 'Resize',
+        # No RESIZE_BICUBIC in BuiltinOperator
+    }
+
+    def __init__(self, TFactory, index):
+        super().__init__(TFactory, index)
+        # Four choices: half_pixel, pytorch_half_pixel, align_corners, asymmetric, tf_crop_and_resize
+        self.attrs['coordinate_transformation_mode'] = 'half_pixel'
+        # This attribute is valid only if "mode" is "cubic".
+        # The coefficient 'a' used in cubic interpolation.
+        # Two common choice are -0.5 (in some cases of TensorFlow) and -0.75 (in PyTorch).
+        self.attrs['cubic_coeff_a'] = -0.75
+        self.attrs['exclude_outside'] = 0
+        self.attrs['extrapolation_value'] = 0.0
+        # Three interpolation modes: nearest (default), linear and cubic.
+        # The "linear" mode includes linear interpolation for 1D tensor and N-linear interpolation for N-D tensor
+        # (for example, bilinear interpolation for 2D tensor).
+        # The "cubic" mode includes cubic interpolation for 1D tensor and N-cubic interpolation for N-D tensor
+        # (for example, bicubic interpolation for 2D tensor).
+        self.attrs['mode'] = 'nearest'
+        # Four modes: round_prefer_floor (default, as known as round half down),
+        # round_prefer_ceil (as known as round half up), floor, ceil.
+        # Only used by nearest interpolation.
+        # It indicates how to get "nearest" pixel in input tensor from x_original,
+        # so this attribute is valid only if "mode" is "nearest".
+        self.attrs['nearest_mode'] = 'round_prefer_floor'
+
+        self.setInited()
+
+    @property
+    def type(self):
+        return 'Resize'
+
+    @property
+    def isRESIZE_BILINEAR(self):
+        op = self.tflite
+        opcode = self.model.OperatorCodes(op.OpcodeIndex()).BuiltinCode()
+        return opcode is tflite.BuiltinOperator.RESIZE_BILINEAR
+
+    @property
+    def isRESIZE_NEAREST_NEIGHBOR(self):
+        op = self.tflite
+        opcode = self.model.OperatorCodes(op.OpcodeIndex()).BuiltinCode()
+        return opcode is tflite.BuiltinOperator.RESIZE_NEAREST_NEIGHBOR
+
+    def parse(self):
+        logger.debug("Parsing %s...", self.type)
+        op = self.tflite
+        opcode = self.model.OperatorCodes(op.OpcodeIndex()).BuiltinCode()
+        assert(opcode is tflite.BuiltinOperator.RESIZE_BILINEAR or tflite.BuiltinOperator.RESIZE_NEAREST_NEIGHBOR)
+
+        assert(op.InputsLength() >= 2), "Only first two arguments: image and size, are compulsory"
+        assert(op.OutputsLength() == 1)
+
+        # image
+        im = self.parseInput(0)
+
+        # Fill 'ROI' empty temporarily
+        # because 'tf_crop_and_resize' was not found in BuiltinOptions of ResizeBilinear and ResizeNearestNeighbor
+        holder = self.TFactory.createVector(im, np.array([]))
+        holder.parse()
+        holder.addConsumer(self)
+        self.inputs.append(holder)
+
+        # Fill empty to scales temporarily
+        sc = self.TFactory.createVector(im, np.array([]))
+        sc.parse()
+        self.inputs.append(sc)
+
+        # output size expected
+        sz = self.parseInput(1)
+
+        # In case Resize, the number of elements of both the arguments
+        # 'sizes' and 'scales' are required to be the same as the rank of input 'X'.
+        # The 'X' usually has a (N,C,H,W) layout after transforming
+        # while the 'sizes' and 'scales' only have (H_new,W_new).
+        # Thus we copy the N and C to achieve (N, C, H_new,W_new) for these two arguments.
+        assert len(sz.data) == 2
+        assert len(im.shape) == 4
+        sz.data = np.concatenate((np.array([im.shape[0], im.shape[-1]]), sz.data))
+        sz.shape = [len(im.shape)]
+        sz.dtype = mapping.DTYPE_NAME2ONNX['int64']
+
+        ot = self.parseOutput(0)
+
+        # options
+        op_opt = op.BuiltinOptions()
+        option = tflite.ResizeBilinearOptions() if self.isRESIZE_BILINEAR else tflite.ResizeNearestNeighborOptions()
+        option.Init(op_opt.Bytes, op_opt.Pos)
+
+        if self.isRESIZE_BILINEAR:
+            self.attrs['mode'] = 'linear'
+        if option.AlignCorners():
+            self.attrs['coordinate_transformation_mode'] = 'align_corners'
+        # Actually docs for the argument `half_pixel_centers` in tf.compat.v1.image.resize_bilinear
+        # and tf.compat.v1.image.resize_nearest_neighbor are missing...
+        # if option.HalfPixelCenters():
+        #     self.attrs['coordinate_transformation_mode'] = 'half_pixel'
+
+        self.setParsed()
+
+    def propagatableTensors(self):
+        return list()
+
+    def transform(self):
+        pass