Merge pull request #76 from onnx/gs/lstm

allow shape override for inputs

Author: guschmue

README.md

@@ -68,7 +68,7 @@ Parameters:
 - inputs/outputs: Tensorflow graph's input/output names, which can be found with [summarize graph tool](#summarize_graph).
 - target: There are different onnx versions and workarounds for runtimes that can be set with ```--target TARGET```. 
 - opset: by default we uses the newest opset installed with the onnx package (for example onnx-1.2.2 would have opset 7). By specifieing ```--opset``` the user can override the default to generate a graph with the desired opset. For example ```--opset 5``` would create a onnx graph that uses only ops available in opset 5. Because older opsets have in most cases fewer ops, some models might not convert on a older opset.
-- custom-ops: the runtime may support custom ops that are not defined in onnx. A user can asked the converter to map to custom ops by listing them with the --custom-ops option. Tensorflow ops listed here will be mapped to a custom op of the same name as the tensorflow op but in the onnx domain ai.onnx.converters.tensorflow. For example: ```--custom-ops Print``` will insert a op ```Print``` in the onnx domain ```ai.onnx.converters.tensorflow``` into the graph. We also support a python api for custom ops documented later in this readme.
+- custom-ops: the runtime may support custom ops that are not defined in onnx. A user can asked the converter to map to custom ops by listing them with the --custom-ops option. Tensorflow ops listed here will be mapped to a custom op of the same name as the tensorflow op but in the onnx domain ai.onnx.converters.tensorflow. For example: ```--custom-ops Print``` will insert a op ```Print``` in the onnx domain ```ai.onnx.converters.tensorflow``` into the graph. We also support a python api for custom ops documented later in this readme. 
 
 Usage example (run following commands in tensorflow-onnx root directory):
 ```
@@ -79,6 +79,8 @@ python -m tf2onnx.convert\
     --output tests/models/fc-layers/model.onnx\
     --verbose
 ```
+Some models specify placeholders with unknown ranks which can not be mapped to onnx. 
+In those cases one can add the shape behind the input name in ```[]```, for example ```--input X:0[1,28,28,3]```
 
 ## <a name="summarize_graph"></a>Tool to get Graph Inputs & Outputs
 

tests/run_pretrained_models.py

@@ -103,7 +103,7 @@ class Test(object):
 
     def __init__(self, url, local, make_input, input_names, output_names,
                  disabled=False, more_inputs=None, rtol=0.01, atol=0.,
-                 check_only_shape=False, model_type="frozen"):
+                 check_only_shape=False, model_type="frozen", force_input_shape=False):
         self.url = url
         self.make_input = make_input
         self.local = local
@@ -118,6 +118,7 @@ def __init__(self, url, local, make_input, input_names, output_names,
         self.tf_runtime = 0
         self.onnx_runtime = 0
         self.model_type = model_type
+        self.force_input_shape = force_input_shape
 
     def download_file(self):
         """Download file from url."""
@@ -171,9 +172,9 @@ def run_tensorflow(self, sess, inputs):
         return result
 
     @staticmethod
-    def to_onnx(tf_graph, opset=None):
+    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)
+        return process_tf_graph(tf_graph, continue_on_error=False, opset=opset, shape_override=shape_override)
 
     def run_caffe2(self, name, onnx_graph, inputs):
         """Run test again caffe2 backend."""
@@ -253,6 +254,8 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
         """Run complete test against backend."""
         print(name)
         self.perf = perf
+
+        # get the model
         if self.url:
             _, dir_name = self.download_file()
             model_path = os.path.join(dir_name, self.local)
@@ -270,6 +273,7 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
                 tf.train.write_graph(frozen_graph, dir_name, "frozen.pb", as_text=False)
             model_path = os.path.join(dir_name, "frozen.pb")
 
+        # create the input data
         inputs = self.make_input(self.input_names)
         if self.more_inputs:
             for k, v in self.more_inputs.items():
@@ -278,8 +282,9 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
         graph_def = graph_pb2.GraphDef()
         with open(model_path, "rb") as f:
             graph_def.ParseFromString(f.read())
-        graph_def = tf2onnx.tfonnx.tf_optimize(None, inputs, self.output_names, graph_def)
 
+        graph_def = tf2onnx.tfonnx.tf_optimize(None, inputs, self.output_names, graph_def)
+        shape_override = {}
         g = tf.import_graph_def(graph_def, name='')
         with tf.Session(graph=g) as sess:
 
@@ -290,12 +295,16 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
                 if type != "float32":
                     v = inputs[k]
                     inputs[k] = v.astype(dtype)
+            if self.force_input_shape:
+                shape_override = self.input_names
 
+            # run the model with tensorflow
             tf_results = self.run_tensorflow(sess, inputs)
             onnx_graph = None
             print("\ttensorflow", "OK")
             try:
-                onnx_graph = self.to_onnx(sess.graph, opset=opset)
+                # convert model to onnx
+                onnx_graph = self.to_onnx(sess.graph, opset=opset, shape_override=shape_override)
                 print("\tto_onnx", "OK")
                 if debug:
                     onnx_graph.dump_graph()
@@ -362,7 +371,7 @@ def tests_from_yaml(fname):
         input_func = v.get("input_get")
         input_func = _INPUT_FUNC_MAPPING[input_func]
         kwargs = {}
-        for kw in ["rtol", "atol", "disabled", "more_inputs", "check_only_shape", "model_type"]:
+        for kw in ["rtol", "atol", "disabled", "more_inputs", "check_only_shape", "model_type", "force_input_shape"]:
             if v.get(kw) is not None:
                 kwargs[kw] = v[kw]
 

tests/run_pretrained_models.yaml

@@ -129,6 +129,7 @@ mobilenet_v2_1.4_224:
   url: https://storage.googleapis.com/mobilenet_v2/checkpoints/mobilenet_v2_1.4_224.tgz
   model: mobilenet_v2_1.4_224_frozen.pb
   input_get: get_beach
+  force_input_shape: true
   inputs:
     "input:0": [1, 224, 224, 3]
   outputs:

tf2onnx/convert.py

@@ -9,6 +9,7 @@
 from __future__ import print_function
 
 import argparse
+import re
 import sys
 
 import onnx
@@ -36,8 +37,22 @@ def get_args():
     parser.add_argument("--verbose", help="verbose output", action="store_true")
     args = parser.parse_args()
 
+    args.shape_override = None
     if args.inputs:
-        args.inputs = args.inputs.split(",")
+        inputs = []
+        shapes = {}
+        # input takes in most cases the format name:0, where 0 is the output number
+        # in some cases placeholders don't have a rank which onnx can't handle so we let uses override the shape
+        # by appending the same, ie : [1,28,28,3]
+        #
+        pattern = r"(?:([\w:]+)(\[[\d,]+\])?),?"
+        splits = re.split(pattern, args.inputs)
+        for i in range(1, len(splits), 3):
+            inputs.append(splits[i])
+            if splits[i+1] is not None:
+                shapes[splits[i]] = [int(n) for n in splits[i+1][1:-1].split(",")]
+        args.inputs = inputs
+        args.shape_override = shapes
     if args.outputs:
         args.outputs = args.outputs.split(",")
     if args.target:
@@ -85,7 +100,8 @@ def main():
                              target=args.target,
                              opset=args.opset,
                              custom_op_handlers=custom_ops,
-                             extra_opset=extra_opset)
+                             extra_opset=extra_opset,
+                             shape_override=args.shape_override)
 
     model_proto = g.make_model(
         "converted from {}".format(args.input), args.inputs, args.outputs,

tf2onnx/tfonnx.py

@@ -34,7 +34,7 @@
 DEFAULT_TARGET = [TARGET_RS4, TARGET_CAFFE2]
 
 
-def tensorflow_to_onnx(graph):
+def tensorflow_to_onnx(graph, shape_override):
     """
     Load tensorflow graph into an onnx graph with minimal rewrites so
     we can use the onnx graph as intermediate graph.
@@ -58,10 +58,12 @@ def tensorflow_to_onnx(graph):
     # create dict with output to shape mappings
     for node in ops:
         for out in node.outputs:
-            try:
-                shape = out.get_shape().as_list()
-            except Exception as ex:
-                shape = []
+            shape = shape_override.get(out.name)
+            if shape is None:
+                try:
+                    shape = out.get_shape().as_list()
+                except Exception as ex:
+                    shape = []
             dtypes[out.name] = utils.map_tf_dtype(out.dtype)
             output_shapes[out.name] = shape
 
@@ -1306,9 +1308,9 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
             onnx_node = func(g, node, node.name, args)
         except Exception as ex:
             type_, value_, traceback_ = sys.exc_info()
-            ex = traceback.format_exception(type_, value_, traceback_)
+            ex_ext = traceback.format_exception(type_, value_, traceback_)
             if continue_on_error:
-                print(ex)
+                print(ex_ext)
                 onnx_nodes.append(node)
             else:
                 raise ex
@@ -1336,7 +1338,8 @@ def tf_optimize(sess, inputs, outputs, graph_def):
 
 
 def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=None,
-                     opset=None, custom_op_handlers=None, custom_rewriter=None, extra_opset=None):
+                     opset=None, custom_op_handlers=None, custom_rewriter=None,
+                     extra_opset=None, shape_override=None):
     """Convert tensorflow graph to onnx graph.
         Args:
             tf_graph: tensorflow graph
@@ -1359,10 +1362,12 @@ def topological_sort(ops):
                 # if we continue on error, ignore graph cycles so we can report all missing ops
                 pass
 
+    if shape_override is None:
+        shape_override = {}
     if target is None:
         target = DEFAULT_TARGET
 
-    onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes = tensorflow_to_onnx(tf_graph)
+    onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes = tensorflow_to_onnx(tf_graph, shape_override)
 
     g = Graph(onnx_nodes, output_shapes, dtypes, target, opset, extra_opset)
     ops = g.get_nodes()