Add target for converting to channels last format (#1581)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tf2onnx/constants.py

@@ -29,7 +29,9 @@
 TARGET_RS6 = "rs6"
 TARGET_CAFFE2 = "caffe2"
 TARGET_TENSORRT = "tensorrt"
-POSSIBLE_TARGETS = [TARGET_RS4, TARGET_RS5, TARGET_RS6, TARGET_CAFFE2, TARGET_TENSORRT]
+TARGET_CHANNELS_LAST = "nhwc"
+TARGET_CHANNELS_FIRST = "nchw"
+POSSIBLE_TARGETS = [TARGET_RS4, TARGET_RS5, TARGET_RS6, TARGET_CAFFE2, TARGET_TENSORRT, TARGET_CHANNELS_LAST]
 DEFAULT_TARGET = []
 
 NCHW_TO_NHWC = [0, 2, 3, 1]

tf2onnx/graph.py

@@ -1124,13 +1124,29 @@ def make_graph(self, doc, graph_name=None, external_tensor_storage=None):
         # create output_tensor_values
         output_tensor_values = self.make_onnx_graph_io(self.outputs)
 
+        tensor_value_info = []
+
+        for op in ops:
+            if op.domain in [constants.ONNX_DOMAIN, constants.AI_ONNX_ML_DOMAIN]:
+                continue
+            # We still don't 100% trust the accuracy of all the shapes in graph.py, but for custom ops they are
+            # almost certainly accurate and onnx has no other way of knowing them.
+            for out in op.output:
+                if out == '' or out in self.outputs:
+                    continue
+                dtype = self.get_dtype(out)
+                shape = self.get_shape(out)
+                v = utils.make_onnx_inputs_outputs(out, dtype, shape)
+                tensor_value_info.append(v)
+
         # create graph proto
         graph = helper.make_graph([op.op for op in ops],
                                   graph_name,
                                   input_tensor_values,
                                   output_tensor_values,
                                   initializer=initializers,
-                                  doc_string=doc)
+                                  doc_string=doc,
+                                  value_info=tensor_value_info)
 
         return graph
 
@@ -1628,10 +1644,11 @@ def get_onnx_model_properties(onnx_model_proto):
         return kwargs
 
     @staticmethod
-    def create_graph_from_onnx_model(onnx_model_proto):
+    def create_graph_from_onnx_model(onnx_model_proto, target=None):
         """Create Graph loading onnx model proto."""
         # apply shape inference on the model
         inferred_model = shape_inference.infer_shapes(onnx_model_proto)
+        utils.initialize_name_counter(inferred_model)
         graph_proto = inferred_model.graph
 
         opset_version = None
@@ -1644,11 +1661,11 @@ def create_graph_from_onnx_model(onnx_model_proto):
                 extra_opset.append(opset)
 
         utils.make_sure(opset_version is not None, "opset version is not specified for onnx domain")
-        main_graph = GraphUtil.create_graph_from_onnx_graph(graph_proto, opset_version, extra_opset)
+        main_graph = GraphUtil.create_graph_from_onnx_graph(graph_proto, opset_version, extra_opset, target)
         return main_graph
 
     @staticmethod
-    def create_graph_from_onnx_graph(graph_proto, opset_version=None, extra_opset=None):
+    def create_graph_from_onnx_graph(graph_proto, opset_version=None, extra_opset=None, target=None):
         """Create Graph loading onnx graph proto."""
         output_shapes = {}
         output_dtypes = {}
@@ -1675,7 +1692,7 @@ def create_graph_from_onnx_graph(graph_proto, opset_version=None, extra_opset=No
         for n in graph_proto.output:
             output_names.append(n.name)
 
-        g = Graph(nodes_to_append, output_shapes, output_dtypes, None, opset_version, extra_opset, None, output_names)
+        g = Graph(nodes_to_append, output_shapes, output_dtypes, target, opset_version, extra_opset, None, output_names)
         const_nodes = GraphUtil._parse_graph_initializer(g, graph_proto)
         GraphUtil._parse_graph_input(g, graph_proto, [n.name for n in const_nodes])
 
@@ -1702,6 +1719,10 @@ def _parse_shape_and_type_from_value_infos(value_infos):
         for shape_info in value_infos:
             type_proto = shape_info.type
             elem_type = type_proto.tensor_type.elem_type
+            output_dtypes[shape_info.name] = elem_type
+            if not type_proto.tensor_type.HasField("shape"):
+                output_shapes[shape_info.name] = None
+                continue
             shape = type_proto.tensor_type.shape
             tuned_shape = []
             for d in shape.dim:
@@ -1713,7 +1734,6 @@ def _parse_shape_and_type_from_value_infos(value_infos):
                     # it is found, some unknown dims is missing after inference.
                     tuned_shape.append(-1)
             output_shapes[shape_info.name] = tuned_shape
-            output_dtypes[shape_info.name] = elem_type
 
         return output_shapes, output_dtypes
 

tf2onnx/late_rewriters/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""tf2onnx.late_rewriters module."""
+
+from tf2onnx.late_rewriters.channel_order_rewriters import rewrite_channels_first, rewrite_channels_last
+
+
+__all__ = [
+    "rewrite_channels_first",
+    "rewrite_channels_last",
+]

tf2onnx/late_rewriters/channel_order_rewriters.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""
+tf2onnx.late_rewriters.channel_order_rewriters - contains rewriters for replacing ops with channel first/last versions
+"""
+
+from tf2onnx import utils, constants
+
+# pylint: disable=invalid-name,unused-argument,missing-docstring, unused-variable
+
+_CHANNELS_FIRST_OPS = [
+    "AveragePool",
+    "BatchNormalization",
+    "Conv",
+    "ConvInteger",
+    "ConvTranspose",
+    "GlobalAveragePool",
+    "GlobalLpPool",
+    "GlobalMaxPool",
+    "InstanceNormalization",
+    "LpPool",
+    "LRN",
+    "MaxPool",
+    "MaxRoiPool",
+    "MaxUnpool",
+    "QLinearConv",
+]
+
+
+def channel_last_to_first_perm(rank):
+    return [0, rank - 1] + list(range(1, rank - 1))
+
+
+def channel_first_to_last_perm(rank):
+    return [0] + list(range(2, rank)) + [1]
+
+
+def _to_channel_last_handler(g, op):
+    # For now, all ops can use the same handlers (input[0] and output[0] are always correct)
+    rank = g.get_rank(op.output[0])
+    utils.make_sure(rank is not None, "Cannot convert %s node %s with unknown rank to channels last", op.type, op.name)
+    op.type = "ChannelsLast" + op.type
+    op.domain = constants.CONTRIB_OPS_DOMAIN
+    inp_perm = channel_first_to_last_perm(rank)
+    out_perm = channel_last_to_first_perm(rank)
+    output_shape = g.get_shape(op.output[0])
+    if output_shape is not None:
+        output_shape = [output_shape[i] for i in inp_perm]
+    g.set_shape(op.output[0], output_shape)
+
+    g.insert_new_node_on_input(op, "Transpose", op.input[0], input_index=0, perm=inp_perm)
+    g.insert_new_node_on_output("Transpose", op.output[0], perm=out_perm)
+
+
+def _to_channel_first_handler(g, op):
+    rank = g.get_rank(op.output[0])
+    utils.make_sure(rank is not None, "Cannot convert %s node %s with unknown rank to channels last", op.type, op.name)
+    op.type = op.type.replace("ChannelsLast", "")
+    op.domain = constants.ONNX_DOMAIN
+    inp_perm = channel_last_to_first_perm(rank)
+    out_perm = channel_first_to_last_perm(rank)
+    output_shape = g.get_shape(op.output[0])
+    if output_shape is not None:
+        output_shape = [output_shape[i] for i in inp_perm]
+    g.set_shape(op.output[0], output_shape)
+
+    g.insert_new_node_on_input(op, "Transpose", op.input[0], input_index=0, perm=inp_perm)
+    g.insert_new_node_on_output("Transpose", op.output[0], perm=out_perm)
+
+
+def get_channels_first_ops(opset=None):
+    # opset doesn't matter for now
+    return set(_CHANNELS_FIRST_OPS)
+
+
+def rewrite_channels_last(g, ops):
+    channel_first_ops = get_channels_first_ops(g.opset)
+    for op in ops:
+        if op.type in channel_first_ops:
+            _to_channel_last_handler(g, op)
+    return g.get_nodes()
+
+
+def rewrite_channels_first(g, ops):
+    for op in ops:
+        if op.type.startswith("ChannelsLast"):
+            _to_channel_first_handler(g, op)
+    return g.get_nodes()

tf2onnx/optimizer/transpose_optimizer.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 import onnx
-from tf2onnx.constants import NCHW_TO_NHWC, NHWC_TO_NCHW, NCDHW_TO_NDHWC, NDHWC_TO_NCDHW
+from tf2onnx.constants import NCHW_TO_NHWC, NHWC_TO_NCHW, NCDHW_TO_NDHWC, NDHWC_TO_NCDHW, TARGET_CHANNELS_LAST
 from .. import utils
 from .optimizer_base import GraphOptimizerBase
 
@@ -362,14 +362,19 @@ def _should_push_transpose(self, trans, node):
         perm = trans.get_attr_value("perm")
         optimization_gains = 0
         removed_nchws = 0
+        perm_to_push_down = [NCHW_TO_NHWC, NCDHW_TO_NDHWC]
+        perm_to_push_up = [NHWC_TO_NCHW, NDHWC_TO_NCDHW]
+        if self._g.is_target(TARGET_CHANNELS_LAST):
+            perm_to_push_down, perm_to_push_up = perm_to_push_up, perm_to_push_down
+
         for n, inp_id in zip(node.inputs, node.input):
             if is_tranpose_of_type(n, perm):
                 optimization_gains += self._cost_to_transpose(n.inputs[0], n.input[0])
-                if perm in [NCHW_TO_NHWC, NCDHW_TO_NDHWC]:
+                if perm in perm_to_push_down:
                     removed_nchws += 1
             else:
                 optimization_gains -= self._cost_to_transpose(n, inp_id)
-                if perm in [NHWC_TO_NCHW, NDHWC_TO_NCDHW]:
+                if perm in perm_to_push_up:
                     removed_nchws -= 1
         if removed_nchws != 0:
             # Always push nchw transposes if possible

tf2onnx/tfonnx.py

@@ -19,6 +19,7 @@
 from tf2onnx.graph import Graph
 from tf2onnx.rewriter import *  # pylint: disable=wildcard-import
 from tf2onnx.tflite_rewriters import *  # pylint: disable=wildcard-import
+from tf2onnx.late_rewriters import rewrite_channels_last
 from tf2onnx.shape_inference import infer_shape
 from tf2onnx.tf_loader import is_function, resolve_functions, set_function
 from tf2onnx.tf_utils import tensorflow_to_onnx, get_tf_version, compute_const_folding_using_tf
@@ -640,6 +641,8 @@ def compat_handler(ctx, node, **kwargs):
         late_rewriters.append(rewrite_incomplete_type_support_rs5)
     if constants.TARGET_RS6 in target:
         late_rewriters.append(rewrite_incomplete_type_support_rs6)
+    if constants.TARGET_CHANNELS_LAST in target:
+        late_rewriters.append(rewrite_channels_last)
     if late_rewriters:
         run_rewriters(g, late_rewriters, continue_on_error)
 

tf2onnx/utils.py

@@ -184,6 +184,35 @@ def find_opset(opset):
     return opset
 
 
+def get_subgraphs_from_onnx(model_proto):
+    """Returns an iterator over the graphs/subgraphs of a model (using dfs)"""
+    stack = [model_proto.graph]
+    while stack:
+        g = stack.pop()
+        yield g
+        for node in g.node:
+            for attr in node.attribute:
+                if hasattr(attr, "g"):
+                    stack.append(attr.g)
+                if hasattr(attr, "graphs"):
+                    stack.extend(attr.graphs)
+
+
+def initialize_name_counter(model_proto):
+    """Avoid name conflicts by initializing the counter used by make_name based on the provided model"""
+    suffix_regex = re.compile(r"__(\d+)(:\d+)?$")
+    def avoid_name(name):
+        global INTERNAL_NAME
+        suffix = suffix_regex.search(name)
+        if suffix:
+            INTERNAL_NAME = max(INTERNAL_NAME, int(suffix.group(1)) + 1)
+    for g in get_subgraphs_from_onnx(model_proto):
+        for n in g.node:
+            avoid_name(n.name)
+            for out in n.output:
+                avoid_name(out)
+
+
 def save_onnx_model(save_path_root, onnx_file_name, feed_dict, model_proto, include_test_data=False, as_text=False,
                     external_tensor_storage=None):
     """Save onnx model as file. Save a pbtxt file as well if as_text is True"""

tools/onnx-optimize.py

@@ -16,7 +16,8 @@
 from onnx import helper
 
 from tf2onnx.graph import GraphUtil
-from tf2onnx import logging, optimizer
+from tf2onnx import logging, optimizer, constants
+from tf2onnx.late_rewriters import rewrite_channels_first, rewrite_channels_last
 
 
 logging.basicConfig(level=logging.INFO)
@@ -28,23 +29,32 @@ def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--input", required=True, help="onnx input model file")
     parser.add_argument("--output", help="output model file")
+    target_options = [constants.TARGET_CHANNELS_LAST, constants.TARGET_CHANNELS_FIRST]
+    parser.add_argument("--target", default=",".join(constants.DEFAULT_TARGET), choices=target_options,
+                        help="target platform")
     args = parser.parse_args()
+    args.target = args.target.split(",")
     return args
 
 
-def load_graph(fname):
+def load_graph(fname, target):
     model_proto = onnx.ModelProto()
     with open(fname, "rb") as f:
         data = f.read()
         model_proto.ParseFromString(data)
-    g = GraphUtil.create_graph_from_onnx_model(model_proto)
+    g = GraphUtil.create_graph_from_onnx_model(model_proto, target)
     return g, model_proto
 
 
 def main():
     args = get_args()
 
-    g, org_model_proto = load_graph(args.input)
+    g, org_model_proto = load_graph(args.input, args.target)
+
+    if g.is_target(constants.TARGET_CHANNELS_FIRST):
+        g.reset_nodes(rewrite_channels_first(g, g.get_nodes()))
+    if g.is_target(constants.TARGET_CHANNELS_LAST):
+        g.reset_nodes(rewrite_channels_last(g, g.get_nodes()))
 
     g = optimizer.optimize_graph(g)