Merge branch 'master' of https://github.com/onnx/tensorflow-onnx into bench

sdpython · sdpython · commit 45dcb15089de · 2021-06-30T11:07:21.000+02:00
diff --git a/README.md b/README.md
@@ -16,10 +16,10 @@ The common issues we run into we try to document here [Troubleshooting Guide](Tr
 
 <br/>
 
-| Build Type | OS | Python | Tensorflow | Onnx opset | Status |
+| Build Type | OS | Python | Tensorflow | ONNX opset | Status |
 | ---        | ---    | ---    | ---        | ---        | ---    |
-| Unit Test - Basic | Linux, MacOS<sup>\*</sup>, Windows<sup>\*</sup> | 3.6, 3.7, 3.8 | 1.12-1.15, 2.1-2.5 | 7-13 | [![Build Status](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_apis/build/status/unit_test?branchName=master)](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_build/latest?definitionId=16&branchName=master) |
-| Unit Test - Full | Linux, MacOS, Windows | 3.6, 3.7, 3.8 | 1.12-1.15, 2.1-2.5 | 7-13 | [![Build Status](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_apis/build/status/unit_test-matrix?branchName=master)](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_build/latest?definitionId=18&branchName=master) | |
+| Unit Test - Basic | Linux, MacOS<sup>\*</sup>, Windows<sup>\*</sup> | 3.6-3.9 | 1.12-1.15, 2.1-2.5 | 8-15 | [![Build Status](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_apis/build/status/unit_test?branchName=master)](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_build/latest?definitionId=16&branchName=master) |
+| Unit Test - Full | Linux, MacOS, Windows | 3.6-3.9 | 1.12-1.15, 2.1-2.5 | 8-14 | [![Build Status](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_apis/build/status/unit_test-matrix?branchName=master)](https://dev.azure.com/tensorflow-onnx/tensorflow-onnx/_build/latest?definitionId=18&branchName=master) | |
 <br/>
 
 ## Supported Versions
@@ -28,21 +28,23 @@ The common issues we run into we try to document here [Troubleshooting Guide](Tr
 
 tf2onnx will use the ONNX version installed on your system and installs the latest ONNX version if none is found.
 
-We support ONNX opset-6 to opset-13. By default we use ```opset-9``` for the resulting ONNX graph since most runtimes will support opset-9.
+We support and test ONNX opset-8 to opset-14. opset-6 and opset-7 should work but we don't test them.
+By default we use ```opset-9``` for the resulting ONNX graph since most runtimes will support opset-9.
 
 If you want the graph to be generated with a specific opset, use ```--opset``` in the command line, for example ```--opset 13```.
 
 ### TensorFlow
 
-We support ```tf-1.x graphs``` and ```tf-2```. To keep our test matrix manageable we test tf2onnx running on top of ```tf-1.12 and up```.
+We support ```tf-1.x graphs``` and ```tf-2```. To keep our test matrix manageable we test tf2onnx running on top of ```tf-1.12 or better```.
 
 When running under tf-2.x tf2onnx will use the tensorflow V2 controlflow.
 
 You can install tf2onnx on top of tf-1.x or tf-2.x.
 
 ### Python
 
-We support Python ```3.6```, ```3.7``` and ```3.8```.
+We support Python ```3.6-3.9```.
+Note that on windows for Python > 3.7 the protobuf package doesn't use the cpp implementation and is very slow - we recommend to use Python 3.7 for that reason.
 
 ## Prerequisites
 
diff --git a/ci_build/azure_pipelines/pretrained_model_test-matrix.yml b/ci_build/azure_pipelines/pretrained_model_test-matrix.yml
@@ -32,7 +32,7 @@ jobs:
   parameters:
     platforms: ['linux', 'windows']
     python_versions: ['3.8']
-    tf_versions: ['2.5.0rc1']
+    tf_versions: ['2.5.0']
     job:
       steps:
       - template: 'pretrained_model_test.yml'
diff --git a/ci_build/azure_pipelines/templates/job_generator.yml b/ci_build/azure_pipelines/templates/job_generator.yml
@@ -5,7 +5,7 @@ parameters:
   python_versions: ['3.7']
   tf_versions: ['']
   onnx_versions: ['']
-  onnx_opsets: ['13', '12', '11', '10', '9', '8', '7']
+  onnx_opsets: ['14', '13', '12', '11', '10', '9', '8']
   onnx_backends: {onnxruntime: ['1.8.0']}
   job: {}
   run_setup: 'True'
diff --git a/ci_build/azure_pipelines/templates/unit_test.yml b/ci_build/azure_pipelines/templates/unit_test.yml
@@ -1,7 +1,7 @@
 # Run unit test
 
 parameters:
-    onnx_opsets: ['13', '12', '11', '10', '9', '8', '7']
+    onnx_opsets: ['14', '13', '12', '11', '10', '9', '8']
     skip_tflite_tests: 'True'
     skip_tf_tests: 'False'
 
diff --git a/ci_build/azure_pipelines/unit_test-matrix.yml b/ci_build/azure_pipelines/unit_test-matrix.yml
@@ -40,7 +40,7 @@ stages:
       parameters:
         platforms: ['linux', 'windows']
         python_versions: ['3.8']
-        tf_versions: ['2.5.0rc1']
+        tf_versions: ['2.5.0']
         onnx_opsets: ['']
         job:
           steps:
diff --git a/ci_build/azure_pipelines/unit_test.yml b/ci_build/azure_pipelines/unit_test.yml
@@ -6,7 +6,7 @@ stages:
     - template: 'templates/job_generator.yml'
       parameters:
         python_versions: ['3.8']
-        tf_versions: ['2.5.0rc1']
+        tf_versions: ['2.5.0']
         onnx_opsets: ['']
         skip_tflite_tests: 'False'
         skip_tf_tests: 'True'
@@ -30,7 +30,7 @@ stages:
     - template: 'templates/job_generator.yml'
       parameters:
         python_versions: ['3.8']
-        tf_versions: ['2.5.0rc1']
+        tf_versions: ['2.5.0']
         onnx_opsets: ['']
         job:
           steps:
diff --git a/tf2onnx/constants.py b/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]
@@ -45,5 +47,5 @@
 # Mapping opset to IR version.
 # Note: opset 7 and opset 8 came out with IR3 but we need IR4 because of PlaceholderWithDefault
 OPSET_TO_IR_VERSION = {
-    1: 3, 2: 3, 3: 3, 4: 3, 5: 3, 6: 3, 7: 4, 8: 4, 9: 4, 10: 5, 11: 6, 12: 7, 13: 7
+    1: 3, 2: 3, 3: 3, 4: 3, 5: 3, 6: 3, 7: 4, 8: 4, 9: 4, 10: 5, 11: 6, 12: 7, 13: 7, 14: 7
 }
diff --git a/tf2onnx/graph.py b/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
 
diff --git a/tf2onnx/late_rewriters/__init__.py b/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",
+]
diff --git a/tf2onnx/late_rewriters/channel_order_rewriters.py b/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()
diff --git a/tf2onnx/optimizer/transpose_optimizer.py b/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
diff --git a/tf2onnx/tfonnx.py b/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)
 
diff --git a/tf2onnx/utils.py b/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"""
diff --git a/tools/onnx-optimize.py b/tools/onnx-optimize.py
