OP: enable Resize/Prelu/Sigmoid, modified Reshape (#25)

* OP: enable Resize/Prelu/Sigmoid, modified Reshape

* add resize-bilinear to test_ops.py

* add test models for new operators

* fix implicit layout for Resize operator

* tests: update prelu test model with random alpha

* doc: update supported operators

* OP: fix code style

* OP: code clean for Resize and Reshape, add test model for resize-nearest-neighbor

* update reshape to rigtht version

* OP: clean code for COLUMNS=73;
LINES=26;
export COLUMNS LINES;

* OP: fix code style

Co-authored-by: v_mmrzhang <v_mmrzhang@tencent.com>

docs/operator-support-status.md

@@ -16,14 +16,17 @@ Please add in A-Z style.
     - [x] Padding
 - [x] Dequantize
 - [x] FullyConnected
+- [x] Logistic (Sigmoid)
 - [x] MaxPooling
 - [x] Mean (reduce)
 - [x] Mul
 - [x] Pad
+- [x] PReLU
 - [x] Quantize
 - [x] ReLU
 - [x] ReLU6
 - [x] Reshape
+- [x] Resize
 - [x] Softmax
 - [x] Split
 - [x] StridedSlice

tests/test_ops.py

@@ -38,6 +38,8 @@ def test_ops_implicit_layout():
         'fullyconnected.float32',
         'fullyconnected-relu6.float32',
         'maxpooling.float32',
+        'resize-bilinear.float32',
+        'resize-nearest-neighbor.float32',
     )
 
     for op in OP_LIST:
@@ -73,6 +75,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

@@ -6,10 +6,45 @@
 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 +75,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 +91,14 @@ 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 expand its dimension by insert 1 to its shape
+            alpha = self.parseInput(1)
+            alpha.shape.insert(0, 1)
+
         self.parseOutput(0)
 
         self.setParsed()

tflite2onnx/op/reshape.py

@@ -32,24 +32,41 @@ 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")
+        # This block has been commented
+        # because the `Reshape` with only one input seems like a special case
+        # haven't manage to reproduce currently
+        # 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.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,128 @@
+import logging
+import tflite
+import numpy as np
+
+from tflite2onnx.layout import Layout
+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 in self.TypeMapping)
+
+        assert(op.InputsLength() == 2), "Only first two arguments: image and size, are compulsory"
+        assert(op.OutputsLength() == 1)
+
+        # image
+        # TFLite requires its `Resize` to be NHWC
+        ilayout = Layout('NHWC', 'NCHW')
+        im = self.parseInput(0, ilayout)
+
+        # Fill 'ROI' empty temporarily
+        # because 'tf_crop_and_resize' was not found in BuiltinOptions
+        # of ResizeBilinear and ResizeNearestNeighbor
+        roi = self.TFactory.createVector(im, np.array([]))
+        roi.addConsumer(self)
+        self.inputs.append(roi)
+
+        # Fill empty to scales temporarily
+        sc = self.TFactory.createVector(im, np.array([]))
+        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']
+
+        # output
+        olayout = Layout('NHWC', 'NCHW')
+        self.parseOutput(0, olayout)
+
+        # 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'
+        elif self.isRESIZE_NEAREST_NEIGHBOR:
+            self.attrs['mode'] = 'nearest'
+
+        if option.AlignCorners():
+            self.attrs['coordinate_transformation_mode'] = 'align_corners'
+        elif option.HalfPixelCenters():
+            self.attrs['coordinate_transformation_mode'] = 'half_pixel'
+
+        self.setParsed()
+
+    def propagatableTensors(self):
+        return list()
+
+    def transform(self):
+        pass