[Topi][Op][PyTorch][Vitas] Fix inconsistent kernel layout conventions for conv2d_transpose (#9336)

* fix a lot of initial tests

* make pytorch tests pass

* lint

* add test

* fix bug with layout transform

* change layouts for conv2d_transpose too

* fix vitis tests

* fix qnn conv2d transpose tests

* fix fake quantization pass

* add todo

* lint

* undo just formatting changes

* remove formatting only change

* remove f2qi for later pr

* more frontend tests fixes

* fix a lot of initial tests

* make pytorch tests pass

* lint

* add test

* fix bug with layout transform

* change layouts for conv2d_transpose too

* fix vitis tests

* fix qnn conv2d transpose tests

* fix fake quantization pass

* add todo

* lint

* undo just formatting changes

* remove formatting only change

* remove f2qi for later pr

* more frontend tests fixes

* jostle

* fix keras

* fix another frontend test

* fix things

* jostle ci

Author: AndrewZhaoLuo

python/tvm/relay/frontend/caffe.py

@@ -21,11 +21,12 @@
 import numpy as np
 import tvm
 from tvm.ir import IRModule
+
+from ... import nd as _nd
 from .. import analysis
 from .. import expr as _expr
 from .. import function as _function
 from .. import op as _op
-from ... import nd as _nd
 from .common import ExprTable
 from .common import infer_shape as _infer_shape
 
@@ -514,14 +515,16 @@ def convert_deconv(self, op):
             weight_shape = [-1, conv_params.num_output, kh, kw]
             weight_value = np.asarray(weight.data, np.float32)
             weight_value = np.reshape(weight_value, weight_shape)
+
+            # weight shape is in relay's IOHW format rn, we need it to be OIHW
+            weight_value = np.transpose(weight_value, [1, 0, 2, 3])
         else:
             raise Exception("No weight value of layer {} in caffemodel".format(op.name))
 
         weight_expr = self.exp_tab.new_const(weight_value, dtype="float32")
         in_expr = self.exp_tab.get_expr(inputs[0])
         out = _op.nn.conv2d_transpose(data=in_expr, weight=weight_expr, **params)
         if bias:
-
             bias_value = np.asarray(bias.data, np.float32)
             bias_expr = self.exp_tab.new_const(bias_value, dtype="float32")
             out = _op.nn.bias_add(out, bias_expr)

python/tvm/relay/frontend/keras.py

@@ -355,11 +355,14 @@ def _convert_convolution(inexpr, keras_layer, etab):
         else:
             kernel_layout = "HWIO"
     else:
-        kernel_layout = "OIHW"
+        if is_deconv:
+            kernel_layout = "IOHW"
+        else:
+            kernel_layout = "OIHW"
 
     if is_deconv:
         kernel_h, kernel_w, n_filters, in_channels = weight.shape
-        if kernel_layout == "OIHW":
+        if kernel_layout == "IOHW":
             weight = weight.transpose([3, 2, 0, 1])
     elif is_depthconv:
         kernel_h, kernel_w, in_channels, depth_mult = weight.shape

python/tvm/relay/frontend/mxnet.py

@@ -18,40 +18,50 @@
 """MXNet symbol frontend."""
 import json
 import math
+
 import numpy as np
 import tvm
-from tvm.ir import IRModule
-
 from tvm import relay
+from tvm.ir import IRModule
 from tvm.topi.utils import get_const_tuple
+
+from ... import nd as _nd
 from .. import analysis
 from .. import expr as _expr
 from .. import function as _function
 from .. import op as _op
 from .. import scope_builder as _scope_builder
-from ... import nd as _nd
-
 from .common import StrAttrsDict
-from .common import infer_type as _infer_type
+from .common import get_name as _get_name
 from .common import infer_shape as _infer_shape
+from .common import infer_type as _infer_type
 from .common import infer_value as _infer_value
-from .common import get_name as _get_name
-from .nnvm_common import _rename, _binop_scalar, _rbinop_scalar, _reduce
-from .nnvm_common import _arg_reduce, _init_op, _softmax_op, _cast
-from .nnvm_common import _clip, _transpose, _upsampling
-from .nnvm_common import _elemwise_sum, _reshape
-from .nnvm_common import _warn_not_used
 from .mxnet_qnn_op_utils import (
-    quantize_mxnet_min_max,
-    quantize_conv_weights_bias_channel_mkldnn_from_var,
-    quantize_conv_bias_mkldnn_from_var,
-    get_conv_mkldnn_requantized_scale_outDtype,
     dequantize_mxnet_min_max,
+    get_conv_mkldnn_requantized_scale_outDtype,
     get_mkldnn_int8_scale,
-    get_mkldnn_uint8_scale,
     get_mkldnn_requantize_scale_outDtype,
+    get_mkldnn_uint8_scale,
+    quantize_conv_bias_mkldnn_from_var,
+    quantize_conv_weights_bias_channel_mkldnn_from_var,
+    quantize_mxnet_min_max,
+)
+from .nnvm_common import (
+    _arg_reduce,
+    _binop_scalar,
+    _cast,
+    _clip,
+    _elemwise_sum,
+    _init_op,
+    _rbinop_scalar,
+    _reduce,
+    _rename,
+    _reshape,
+    _softmax_op,
+    _transpose,
+    _upsampling,
+    _warn_not_used,
 )
-
 
 __all__ = ["from_mxnet"]
 
@@ -329,7 +339,7 @@ def _mx_conv2d_transpose(inputs, attrs):
     if "kernel_layout" in attrs.attrs:
         kernel_layout = attrs.get_str("kernel_layout")
     else:
-        kernel_layout = "HWIO" if data_layout == "NHWC" else "OIHW"
+        kernel_layout = "HWIO" if data_layout == "NHWC" else "IOHW"
 
     new_attrs = {}
     new_attrs["channels"] = attrs.get_int("num_filter")

python/tvm/relay/frontend/pytorch.py

@@ -19,8 +19,8 @@
 # pylint: disable=import-outside-toplevel, simplifiable-if-expression, cell-var-from-loop, unnecessary-lambda
 # pylint: disable=missing-function-docstring
 """PT: PyTorch frontend."""
-import itertools
 import functools
+import itertools
 import logging
 import math
 import sys
@@ -40,11 +40,11 @@
 from ..prelude import Prelude, StaticTensorArrayOps
 from ..ty import Any, TensorType, TupleType
 from . import qnn_torch
-from .common import AttrCvt, get_relay_op, unbind, lstm_cell, gru_cell
-from .common import infer_value as _infer_value
+from .common import AttrCvt, get_relay_op, gru_cell
 from .common import infer_shape as _infer_shape
+from .common import infer_value as _infer_value
 from .common import infer_value_simulated as _infer_value_simulated
-from .common import try_infer_value
+from .common import lstm_cell, try_infer_value, unbind
 from .pytorch_utils import is_version_greater_than
 
 __all__ = ["from_pytorch"]
@@ -1010,6 +1010,9 @@ def convolution(self, inputs, input_types):
         elif len(kernel_size) == 2:
             data_layout = "NCHW"
             kernel_layout = "OIHW"
+            if use_transpose:
+                # Transposed convolutions have IOHW layout.
+                kernel_layout = "IOHW"
         else:
             data_layout = "NCW"
             kernel_layout = "OIW"

python/tvm/relay/frontend/qnn_torch.py

@@ -19,7 +19,6 @@
 import logging
 
 import numpy as np
-
 import tvm
 from tvm import relay
 from tvm.relay import expr as _expr
@@ -1043,11 +1042,8 @@ def _impl(inputs, _):
 
         weight_shape = list(infer_shape(weight))
 
-        # Swap I and O dims to match shape relay expects for OIHW
-        weight_shape[0], weight_shape[1] = weight_shape[1], weight_shape[0]
-
         kernel_size = (weight_shape[2], weight_shape[3])
-        out_channels = weight_shape[0]
+        out_channels = weight_shape[1]
 
         conv_out = relay.qnn.op.conv2d_transpose(
             inputs[0],
@@ -1064,7 +1060,7 @@ def _impl(inputs, _):
             channels=out_channels,
             output_padding=output_padding,
             out_dtype="int32",
-            kernel_layout="OIHW",
+            kernel_layout="IOHW",
         )
 
         return _do_bias_and_requantize(

python/tvm/relay/frontend/tensorflow_ops.py

@@ -461,8 +461,11 @@ def _impl(inputs, attr, params, mod):
             raise tvm.error.OpAttributeInvalid(msg.format(attr["padding"]))
 
         if "kernel_layout" not in attr:
-            if opname in ["conv", "conv_transpose"]:
+            if opname == "conv":
                 attr["kernel_layout"] = "HWIO" if attr["data_format"] == "NHWC" else "OIHW"
+            elif opname == "conv_transpose":
+                # conv_transpose in TVM has weights be IOHW for NCHW
+                attr["kernel_layout"] = "HWIO" if attr["data_format"] == "NHWC" else "IOHW"
             else:
                 attr["kernel_layout"] = "HWOI" if attr["data_format"] == "NHWC" else "OIHW"
 

python/tvm/relay/frontend/tflite.py

@@ -16,24 +16,25 @@
 # under the License.
 # pylint: disable=invalid-name, unused-argument, too-many-lines, import-outside-toplevel
 """Tensorflow lite frontend."""
-import math
 import itertools
+import math
+
 import numpy as np
 import tvm
+from tvm import relay
 from tvm.ir import IRModule
 
-from tvm import relay
+from ... import nd as _nd
 from .. import analysis
 from .. import expr as _expr
 from .. import function as _function
 from .. import op as _op
 from .. import qnn as _qnn
-from ... import nd as _nd
 from .common import ExprTable
-from .common import infer_shape as _infer_shape, to_int_list
+from .common import infer_shape as _infer_shape
+from .common import to_int_list
 from .tflite_flexbuffer import FlexBufferDecoder
 
-
 __all__ = ["from_tflite"]
 
 
@@ -53,9 +54,9 @@ class OperatorConverter(object):
     def __init__(self, model, subgraph, exp_tab):
 
         try:
+            from tflite.ActivationFunctionType import ActivationFunctionType
             from tflite.BuiltinOperator import BuiltinOperator
             from tflite.BuiltinOptions import BuiltinOptions
-            from tflite.ActivationFunctionType import ActivationFunctionType
         except ImportError:
             raise ImportError("The tflite package must be installed")
 
@@ -1061,8 +1062,8 @@ def convert_log_softmax(self, op):
     def convert_concatenation(self, op):
         """Convert TFLite concatenation"""
         try:
-            from tflite.ConcatenationOptions import ConcatenationOptions
             from tflite.BuiltinOptions import BuiltinOptions
+            from tflite.ConcatenationOptions import ConcatenationOptions
         except ImportError:
             raise ImportError("The tflite package must be installed")
 
@@ -1242,10 +1243,10 @@ def _convert_elemwise(self, relay_op, op, ignore_qnn_params=False):
         """Generic method to Convert TFLite elemwise"""
         try:
             from tflite.AddOptions import AddOptions
-            from tflite.SubOptions import SubOptions
-            from tflite.MulOptions import MulOptions
-            from tflite.DivOptions import DivOptions
             from tflite.BuiltinOptions import BuiltinOptions
+            from tflite.DivOptions import DivOptions
+            from tflite.MulOptions import MulOptions
+            from tflite.SubOptions import SubOptions
         except ImportError:
             raise ImportError("The tflite package must be installed")
 
@@ -1804,9 +1805,9 @@ def convert_reduce_any(self, op):
     def _convert_arg_min_max(self, relay_op, op):
         """Generic method converting TFLite arg_min_max"""
         try:
-            from tflite.BuiltinOptions import BuiltinOptions
-            from tflite.ArgMinOptions import ArgMinOptions
             from tflite.ArgMaxOptions import ArgMaxOptions
+            from tflite.ArgMinOptions import ArgMinOptions
+            from tflite.BuiltinOptions import BuiltinOptions
         except ImportError:
             raise ImportError("The tflite package must be installed")
 
@@ -1853,8 +1854,8 @@ def convert_arg_max(self, op):
     def convert_fully_connected(self, op):
         """Convert TFLite fully connected"""
         try:
-            from tflite.FullyConnectedOptions import FullyConnectedOptions
             from tflite.BuiltinOptions import BuiltinOptions
+            from tflite.FullyConnectedOptions import FullyConnectedOptions
             from tflite.TensorType import TensorType
         except ImportError:
             raise ImportError("The tflite package must be installed")
@@ -2024,10 +2025,10 @@ def convert_conv(self, op, conv_type):
         """convolution implementation."""
         try:
             from tflite.BuiltinOptions import BuiltinOptions
-            from tflite.TensorType import TensorType
             from tflite.Conv2DOptions import Conv2DOptions
             from tflite.DepthwiseConv2DOptions import DepthwiseConv2DOptions
             from tflite.Padding import Padding
+            from tflite.TensorType import TensorType
         except ImportError:
             raise ImportError("The tflite package must be installed")
 
@@ -2434,8 +2435,8 @@ def convert_pool2d(self, op, pool_type):
         """pool2d implementation."""
         try:
             from tflite.BuiltinOptions import BuiltinOptions
-            from tflite.Pool2DOptions import Pool2DOptions
             from tflite.Padding import Padding
+            from tflite.Pool2DOptions import Pool2DOptions
         except ImportError:
             raise ImportError("The tflite package must be installed")
 
@@ -2850,9 +2851,9 @@ def convert_transpose_conv(self, op):
         """Convert TFLite TRANSPOSE_CONV"""
         try:
             from tflite.BuiltinOptions import BuiltinOptions
+            from tflite.Padding import Padding
             from tflite.TensorType import TensorType
             from tflite.TransposeConvOptions import TransposeConvOptions
-            from tflite.Padding import Padding
         except ImportError:
             raise ImportError("The tflite package must be installed")
 
@@ -2946,7 +2947,7 @@ def convert_transpose_conv(self, op):
                 channels=int(out_channels),
                 kernel_size=(int(kernel_h), int(kernel_w)),
                 data_layout="NHWC",
-                kernel_layout="OIHW",
+                kernel_layout="IOHW",
                 out_dtype="int32",
             )
         else:
@@ -2958,7 +2959,7 @@ def convert_transpose_conv(self, op):
                 channels=int(out_channels),
                 kernel_size=(int(kernel_h), int(kernel_w)),
                 data_layout="NHWC",
-                kernel_layout="OIHW",
+                kernel_layout="IOHW",
                 out_dtype=output_tensor_type_str,
             )
 
@@ -3717,8 +3718,8 @@ def from_tflite(model, shape_dict=None, dtype_dict=None, op_converter=OperatorCo
         The parameter dict to be used by relay
     """
     try:
-        import tflite.SubGraph
         import tflite.BuiltinOperator
+        import tflite.SubGraph
     except ImportError:
         raise ImportError("The tflite package must be installed")
 

python/tvm/relay/op/nn/nn.py

@@ -522,7 +522,7 @@ def conv2d_transpose(
     channels=None,
     kernel_size=None,
     data_layout="NCHW",
-    kernel_layout="OIHW",
+    kernel_layout="IOHW",
     out_layout="",
     output_padding=(0, 0),
     out_dtype="",

python/tvm/relay/qnn/op/layout_conversions.py

@@ -119,7 +119,7 @@ def convert_qnn_conv2d_transpose(attrs, inputs, tinfos, desired_layouts):
 
     # Handle default kernel layouts
     if desired_data_layout == "NCHW":
-        new_attrs["kernel_layout"] = "OIHW"
+        new_attrs["kernel_layout"] = "IOHW"
         return relay.qnn.op.conv2d_transpose(*inputs, **new_attrs)
     if desired_data_layout == "NHWC":
         new_attrs["kernel_layout"] = "HWIO"