add transpose optimizer, and integrate it after tf graph conversion

tests/run_pretrained_models.py

@@ -21,6 +21,8 @@
 from tensorflow.core.framework import graph_pb2
 from tf2onnx.tfonnx import process_tf_graph
 from tensorflow.python.framework.graph_util import convert_variables_to_constants
+from tf2onnx.optimizer.onnx_graph import OnnxGraph
+from tf2onnx.optimizer.transpose_optimizer import TransposeOptimizer
 
 TMPPATH = tempfile.mkdtemp()
 PERFITER = 1000
@@ -176,10 +178,9 @@ def to_onnx(tf_graph, opset=None, shape_override=None):
         """Convert graph to tensorflow."""
         return process_tf_graph(tf_graph, continue_on_error=False, opset=opset, shape_override=shape_override)
 
-    def run_caffe2(self, name, onnx_graph, inputs):
+    def run_caffe2(self, name, model_proto, inputs):
         """Run test again caffe2 backend."""
         import caffe2.python.onnx.backend
-        model_proto = onnx_graph.make_model("test", inputs.keys(), self.output_names)
         prepared_backend = caffe2.python.onnx.backend.prepare(model_proto)
         results = prepared_backend.run(inputs)
         if self.perf:
@@ -189,12 +190,11 @@ def run_caffe2(self, name, onnx_graph, inputs):
             self.onnx_runtime = time.time() - start
         return results
 
-    def run_onnxmsrt(self, name, onnx_graph, inputs):
+    def run_onnxmsrt(self, name, model_proto, inputs):
         """Run test against onnxmsrt backend."""
         import lotus
         # create model and datafile in tmp path.
         model_path = os.path.join(TMPPATH, name + "_model.pb")
-        model_proto = onnx_graph.make_model("test", inputs.keys(), self.output_names)
         with open(model_path, "wb") as f:
             f.write(model_proto.SerializeToString())
         m = lotus.ModelExecutor(model_path)
@@ -206,11 +206,10 @@ def run_onnxmsrt(self, name, onnx_graph, inputs):
             self.onnx_runtime = time.time() - start
         return results
 
-    def run_onnxmsrtnext(self, name, onnx_graph, inputs):
+    def run_onnxmsrtnext(self, name, model_proto, inputs):
         """Run test against msrt-next backend."""
         import lotus
         model_path = os.path.join(TMPPATH, name + ".pb")
-        model_proto = onnx_graph.make_model("test", inputs.keys(), self.output_names)
         with open(model_path, "wb") as f:
             f.write(model_proto.SerializeToString())
         m = lotus.InferenceSession(model_path)
@@ -222,11 +221,10 @@ def run_onnxmsrtnext(self, name, onnx_graph, inputs):
             self.onnx_runtime = time.time() - start
         return results
 
-    def run_onnxcntk(self, name, onnx_graph, inputs):
+    def run_onnxcntk(self, name, model_proto, inputs):
         """Run test against cntk backend."""
         import cntk as C
         model_path = os.path.join(TMPPATH, name + "_model.pb")
-        model_proto = onnx_graph.make_model("test", inputs.keys(), self.output_names)
         with open(model_path, "wb") as f:
             f.write(model_proto.SerializeToString())
         z = C.Function.load(model_path, format=C.ModelFormat.ONNX)
@@ -242,10 +240,9 @@ def run_onnxcntk(self, name, onnx_graph, inputs):
             self.onnx_runtime = time.time() - start
         return results
 
-    def create_onnx_file(self, name, onnx_graph, inputs, outdir):
+    def create_onnx_file(self, name, model_proto, inputs, outdir):
         os.makedirs(outdir, exist_ok=True)
         model_path = os.path.join(outdir, name + ".onnx")
-        model_proto = onnx_graph.make_model(name, inputs.keys(), self.output_names)
         with open(model_path, "wb") as f:
             f.write(model_proto.SerializeToString())
         print("\tcreated", model_path)
@@ -300,29 +297,33 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
 
             # run the model with tensorflow
             tf_results = self.run_tensorflow(sess, inputs)
-            onnx_graph = None
+            opt_model_proto = None
             print("\ttensorflow", "OK")
             try:
                 # convert model to onnx
                 onnx_graph = self.to_onnx(sess.graph, opset=opset, shape_override=shape_override)
+                model_proto = onnx_graph.make_model("test", inputs.keys(), self.output_names)
+                # optimize the onnx graph with TransposeOptimizer
+                optimizer = TransposeOptimizer(OnnxGraph(model_proto.graph))
+                opt_model_proto = optimizer.optimize()
                 print("\tto_onnx", "OK")
                 if debug:
-                    onnx_graph.dump_graph()
+                    opt_model_proto.dump_graph()
                 if onnx_file:
-                    self.create_onnx_file(name, onnx_graph, inputs, onnx_file)
+                    self.create_onnx_file(name, opt_model_proto, inputs, onnx_file)
             except Exception as ex:
                 print("\tto_onnx", "FAIL", ex)
 
         try:
             onnx_results = None
             if backend == "caffe2":
-                onnx_results = self.run_caffe2(name, onnx_graph, inputs)
+                onnx_results = self.run_caffe2(name, opt_model_proto, inputs)
             elif backend == "onnxmsrt":
-                onnx_results = self.run_onnxmsrt(name, onnx_graph, inputs)
+                onnx_results = self.run_onnxmsrt(name, opt_model_proto, inputs)
             elif backend == "onnxmsrtnext":
-                onnx_results = self.run_onnxmsrtnext(name, onnx_graph, inputs)
+                onnx_results = self.run_onnxmsrtnext(name, opt_model_proto, inputs)
             elif backend == "cntk":
-                onnx_results = self.run_onnxcntk(name, onnx_graph, inputs)
+                onnx_results = self.run_onnxcntk(name, opt_model_proto, inputs)
             else:
                 raise ValueError("unknown backend")
             print("\trun_onnx OK")

tf2onnx/convert.py

@@ -17,7 +17,8 @@
 import tf2onnx.utils
 from tf2onnx.tfonnx import process_tf_graph, tf_optimize, DEFAULT_TARGET, POSSIBLE_TARGETS
 from onnx import helper
-
+from tf2onnx.optimizer.onnx_graph import OnnxGraph
+from tf2onnx.optimizer.transpose_optimizer import TransposeOptimizer
 
 _TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
 
@@ -94,10 +95,14 @@ def main():
         "converted from {}".format(args.input), args.inputs, args.outputs,
         optimize=not args.continue_on_error)
 
+    onnx_graph = OnnxGraph(model_proto.graph)
+    optimizer = TransposeOptimizer(onnx_graph)
+    opt_model_proto = optimizer.optimize()
+
     # write onnx graph
     if args.output:
         with open(args.output, "wb") as f:
-            f.write(model_proto.SerializeToString())
+            f.write(opt_model_proto.SerializeToString())
 
 
 main()

tf2onnx/optimizer/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+__all__ = ["onnx_graph", "transpose_optimizer"]
+
+#import tf2onnx
+#from tf2onnx import tfonnx, utils, graph, graph_matcher