Merge pull request #102 from onnx/gs/lstm

fix for bug 101

Author: guschmue

tests/test_internals.py

@@ -163,6 +163,14 @@ def test_match_flipped(self):
         match_results = list(matcher.match_ops(ops))
         self.assertEqual(1, len(match_results))
 
+    def test_cmdarg_parse(self):
+        arg = "input/V-1_2:0,input/X:0[1,2,3],Y:1[4,5],Z:3,A:1,B"
+        expected_inputs = ['input/V-1_2:0', 'input/X:0', 'Y:1', 'Z:3', 'A:1', 'B']
+        expected_shape = {'Y:1': [4, 5], 'input/X:0': [1, 2, 3]}
+        inputs, shape_override = tf2onnx.utils.split_nodename_and_shape(arg)
+        self.assertEqual(expected_inputs, inputs)
+        self.assertEqual(expected_shape, shape_override)
+
 
 if __name__ == '__main__':
     unittest.main()

tf2onnx/convert.py

@@ -39,20 +39,7 @@ def get_args():
 
     args.shape_override = None
     if args.inputs:
-        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
+        args.inputs, args.shape_override = tf2onnx.utils.split_nodename_and_shape(args.inputs)
     if args.outputs:
         args.outputs = args.outputs.split(",")
     if args.target:

tf2onnx/utils.py

@@ -8,11 +8,13 @@
 from __future__ import division
 from __future__ import print_function
 
+import re
 import numpy as np
 import tensorflow as tf
 from onnx import helper, onnx_pb
 from tensorflow.core.framework import types_pb2, tensor_pb2
 
+
 #
 #  mapping dtypes from tensorflow to onnx
 #
@@ -94,6 +96,24 @@ def make_name(name):
     return "{}__{}".format(name, INTERNAL_NAME)
 
 
+def split_nodename_and_shape(name):
+    # pattern for a node name
+    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]
+    name_pattern = r"(?:([\w\d/\-_:]+)(\[[\d,]+\])?),?"
+    splits = re.split(name_pattern, name)
+    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(",")]
+    if len(shapes) == 0:
+        shapes = None
+    return inputs, shapes
+
+
 def tf_to_onnx_tensor(tensor, name=""):
     """Convert tensorflow tensor to onnx tensor."""
     new_type = TF_TO_ONNX_DTYPE[tensor.dtype]