diff --git a/onnx_tf/handlers/backend/concat.py b/onnx_tf/handlers/backend/concat.py index 2f2a870d..bbd6866f 100644 --- a/onnx_tf/handlers/backend/concat.py +++ b/onnx_tf/handlers/backend/concat.py @@ -3,7 +3,8 @@ from onnx_tf.handlers.backend_handler import BackendHandler from onnx_tf.handlers.handler import onnx_op from onnx_tf.handlers.handler import tf_func - +from onnx_tf.common import get_data_format +from onnx_tf.common import get_perm_from_formats @onnx_op("Concat") @tf_func(tf.concat) @@ -12,7 +13,24 @@ class Concat(BackendHandler): @classmethod def _common(cls, node, **kwargs): inputs = [kwargs["tensor_dict"][inp] for inp in node.inputs] - return [cls.make_tensor_from_onnx_node(node, inputs=[inputs])] + rank = len(inputs[0].get_shape()) + storage_format, compute_format = get_data_format(rank) if (rank >= 2 and rank <= 5) else ('', '') + if storage_format == compute_format: + return [cls.make_tensor_from_onnx_node(node, inputs=[inputs])] + else: + # Transpose from storage_format to compute_format and do concat. + # NOTE: it's assumed that all other operators will be run in `compute_format` as much as + # possible, so those redundant `transpose` operators can be resolved. + inputs = [tf.transpose(x, get_perm_from_formats(storage_format, compute_format)) for x in inputs] + # adjust concat axis according to source and target layout format. + from copy import deepcopy + attrs = deepcopy(node.attrs) + axis = attrs["axis"] + axis = compute_format.index(storage_format[axis]) + attrs["axis"] = axis + output = cls.make_tensor_from_onnx_node(node, inputs=[inputs], attrs=attrs) + output = [tf.transpose(output, get_perm_from_formats(compute_format, storage_format))] + return output @classmethod def version_1(cls, node, **kwargs): diff --git a/onnx_tf/handlers/backend/conv_mixin.py b/onnx_tf/handlers/backend/conv_mixin.py index 76a93153..ede6bb7f 100644 --- a/onnx_tf/handlers/backend/conv_mixin.py +++ b/onnx_tf/handlers/backend/conv_mixin.py @@ -62,34 +62,6 @@ def conv(cls, node, input_dict, transpose=False): pads = node.attrs.get("pads", [0, 0] * spatial_size) - # Check auto_pad nonexistent or NOTSET first - if "auto_pad" not in node.attrs or node.attrs["auto_pad"] == "NOTSET": - if not transpose: - if pads != [0, 0] * spatial_size: - x = PadMixin.get_padding_as_op(x, pads) - pad_mode = "VALID" - else: - pad_mode = "NOTSET" - # Then we use auto_pad to setup pad_mode - elif node.attrs["auto_pad"] == "SAME_UPPER": - pad_mode = "SAME" - elif node.attrs["auto_pad"] == "VALID": - pad_mode = "VALID" - elif node.attrs["auto_pad"] == "SAME_LOWER": - pad_mode = PAD_TF_INCOMPATIBLE - else: - raise ValueError("Invalid auto_pad attribute: {}".format( - node.attrs["auto_pad"])) - - # Currently auto_pad = SAME_LOWER is not supported - if pad_mode is PAD_TF_INCOMPATIBLE: - if transpose: - exception.OP_UNSUPPORTED_EXCEPT( - "ConvTranspose with auto_pad `SAME_LOWER`", "Tensorflow") - else: - exception.OP_UNSUPPORTED_EXCEPT("Conv with auto_pad `SAME_LOWER`", - "Tensorflow") - group = node.attrs.get("group", 1) weight_shape = weights.get_shape().as_list() # Is this convolution depthwise we can support? @@ -134,6 +106,34 @@ def conv(cls, node, input_dict, transpose=False): else: xs = tf.split(x, num_or_size_splits=group, axis=-1) + # Check auto_pad nonexistent or NOTSET first + if "auto_pad" not in node.attrs or node.attrs["auto_pad"] == "NOTSET": + if not transpose: + if pads != [0, 0] * spatial_size: + xs = [PadMixin.get_padding_as_op(x, pads, format=compute_format) for x in xs] + pad_mode = "VALID" + else: + pad_mode = "NOTSET" + # Then we use auto_pad to setup pad_mode + elif node.attrs["auto_pad"] == "SAME_UPPER": + pad_mode = "SAME" + elif node.attrs["auto_pad"] == "VALID": + pad_mode = "VALID" + elif node.attrs["auto_pad"] == "SAME_LOWER": + pad_mode = PAD_TF_INCOMPATIBLE + else: + raise ValueError("Invalid auto_pad attribute: {}".format( + node.attrs["auto_pad"])) + + # Currently auto_pad = SAME_LOWER is not supported + if pad_mode is PAD_TF_INCOMPATIBLE: + if transpose: + exception.OP_UNSUPPORTED_EXCEPT( + "ConvTranspose with auto_pad `SAME_LOWER`", "Tensorflow") + else: + exception.OP_UNSUPPORTED_EXCEPT("Conv with auto_pad `SAME_LOWER`", + "Tensorflow") + if transpose: if dilations != [1] * spatial_size: raise RuntimeError("Cannot set non-1 dilation for conv transpose.") diff --git a/onnx_tf/handlers/backend/pad_mixin.py b/onnx_tf/handlers/backend/pad_mixin.py index 271f4f96..a465ee09 100644 --- a/onnx_tf/handlers/backend/pad_mixin.py +++ b/onnx_tf/handlers/backend/pad_mixin.py @@ -5,13 +5,23 @@ class PadMixin(object): @classmethod - def get_padding_as_op(cls, x, pads): + def get_padding_as_op(cls, x, pads, format:str=None): num_dim = int(len(pads) / 2) - tf_pads = np.transpose(np.array(pads).reshape([2, num_dim])) - tf_pads = [0, 0, 0, 0] + tf_pads.flatten().tolist() + # tf_pads = np.transpose(np.array(pads).reshape([2, num_dim])) + if format is None: + NIdx, CIdx = 0, 1 + else: + assert "N" in format and "C" in format, "expected `N` and `C` in padding op's input format " \ + "if given" + NIdx = format.index("N") + CIdx = format.index("C") + # create an empty tf_pads array + tf_pads = np.zeros([num_dim + 2, 2], dtype=np.int32) + # the indices of spatial axes in input format. + spatial_indices = [axis for axis in range(num_dim + 2) if axis not in [NIdx, CIdx]] + # fill pads into tf_pads's spatial axes + tf_pads[spatial_indices, :] = np.transpose(np.array(pads).reshape([2, num_dim])) - padding = tf.constant( - np.array(tf_pads).reshape([num_dim + 2, 2]) - .astype(np.int32)) # tf requires int32 paddings - return tf.pad(x, padding) + padding = tf.constant(tf_pads) # tf requires int32 paddings + return tf.pad(x, padding) \ No newline at end of file