Support NHWC compute format in concat and conv's pad op.

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):

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.")

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)