Merge pull request #37 from onnx/gs/onnx-1.2

added support for topk, stridedslice(limited), custom opset

Author: guschmue

tests/test_backend.py

@@ -84,7 +84,6 @@ def run_onnxmsrt(onnx_graph, inputs, output_names, test_name):
         model_path = os.path.join(TMPPATH, test_name + ".pb")
         with open(model_path, "wb") as f:
             f.write(onnx_graph.SerializeToString())
-
         m = lotus.ModelExecutor(model_path)
         results = m.run(output_names, inputs)
         return results[0]
@@ -96,7 +95,6 @@ def run_onnxmsrtnext(onnx_graph, inputs, output_names, test_name):
         model_path = os.path.join(TMPPATH, test_name + ".pb")
         with open(model_path, "wb") as f:
             f.write(onnx_graph.SerializeToString())
-
         m = lotus.InferenceSession(model_path)
         results = m.run(output_names, inputs)
         return results[0]
@@ -764,19 +762,32 @@ def test_cancel_transpose(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
-    @unittest.skip
-    def test_strided_slice0(self):
-        # FIXME: not implemented yet
-        x_val = np.array([
-            [[1, 1, 1], [2, 2, 2]],
-            [[3, 3, 3], [4, 4, 4]],
-            [[5, 5, 5], [6, 6, 6]]], dtype=np.float32)
+    def test_topk(self):
+        x_val = np.arange(3*2*3).astype("float32")
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        values, indices = tf.nn.top_k(x, 5, sorted=True)
+        output  = tf.identity(values, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
+    @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "multiple dims not supported")
+    def test_strided_slice1(self):
+        x_val = np.arange(3*2*3).astype("float32").reshape(3, 2, 3)
         x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
         x_ = tf.strided_slice(x, [1, 0, 0], [2, 1, 3], [1, 1, 1])
         output = tf.identity(x_, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
+    @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "multiple dims not supported")
+    def test_strided_slice2(self):
+        x_val = np.arange(3*2*3).astype("float32").reshape(3, 2, 3)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = tf.strided_slice(x, [1, 0, 0], [2, 2, 3], [1, 1, 1])
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
     @unittest.skip
     def test_resize_nearest_neighbor(self):
         # this should work but no runtime I tried supports it.

tests/test_graph.py

@@ -266,14 +266,17 @@ def print_handler(ctx, node, name, args):
             # becomes:
             #   T output = Identity(T Input)
             node.type = "Identity"
+            node.domain = "tf"
             del node.input[1:]
             return node
 
         with tf.Session() as sess:
             x = tf.placeholder(tf.float32, [2, 3], name="input1")
             x_ = tf.Print(x, [x], "hello")
             _ = tf.identity(x_, name="output")
-            g = process_tf_graph(sess.graph, custom_op_handlers={"Print": print_handler})
+            g = process_tf_graph(sess.graph,
+                                 custom_op_handlers={"Print": print_handler},
+                                 extra_opset=helper.make_opsetid("tf", 1))
             self.assertEqual(
                 'digraph { Print [op_type=Identity] output [op_type=Identity] input1:0 -> Print Print:0 -> output }',
                 onnx_to_graphviz(g))

tf2onnx/convert.py

@@ -12,6 +12,7 @@
 import tensorflow as tf
 import tf2onnx.utils
 from tf2onnx.tfonnx import process_tf_graph, tf_optimize, DEFAULT_TARGET, POSSIBLE_TARGETS
+from onnx import helper
 
 
 def get_args():
@@ -44,7 +45,7 @@ def get_args():
 
 
 def default_custom_op_handler(ctx, node, name, args):
-    node.type = "tf." + node.type
+    node.domain = "tf"
     return node
 
 
@@ -57,7 +58,13 @@ def main():
     # support unknown dimensions.
     tf2onnx.utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
 
-    custom_ops = {op: default_custom_op_handler for op in args.custom_ops.split(",")} if args.custom_ops else {}
+    if args.custom_ops:
+        # default custom ops for tensorflow-onnx are in the "tf" namespace
+        custom_ops = {op: default_custom_op_handler for op in args.custom_ops.split(",")}
+        extra_opset = [helper.make_opsetid("tf", 1)]
+    else:
+        args.custom_ops = {}
+        extra_opset = None
 
     graph_def = tf.GraphDef()
     with tf.gfile.FastGFile(args.input, 'rb') as f:
@@ -71,7 +78,8 @@ def main():
                              verbose=args.verbose,
                              target=args.target,
                              opset=args.opset,
-                             custom_op_handlers=custom_ops)
+                             custom_op_handlers=custom_ops,
+                             extra_opset=extra_opset)
 
     model_proto = g.make_model(
         "converted from {}".format(args.input), args.inputs, args.outputs,

tf2onnx/graph.py

@@ -81,6 +81,16 @@ def type(self, val):
         """Set Op type."""
         self._op.op_type = val
 
+    @property
+    def domain(self):
+        """Return Op type."""
+        return self._op.domain
+
+    @type.setter
+    def domain(self, val):
+        """Set Op type."""
+        self._op.domain = val
+
     def is_nhwc(self):
         """Return True if node is in NCHW format."""
         return self.data_format == "NHWC"
@@ -200,7 +210,7 @@ def update_proto(self):
 class Graph(object):
     """"Class that provides graph manipulation and matching."""
 
-    def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=None):
+    def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=None, extra_opset=None):
         """Create Graph.
         Args:
             nodes: list of Node()
@@ -219,9 +229,10 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
         self._output_shapes = output_shapes
         ops = [Node(node, self) for node in nodes]
         self.set_nodes(ops)
-        if opset is None:
+        if opset is None or opset == 0:
             opset = defs.onnx_opset_version()
         self._opset = opset
+        self._extra_opset = extra_opset
 
     @property
     def opset(self):
@@ -401,10 +412,13 @@ def make_model(self, doc, input_names, output_names, optimize=True):
 
         kwargs = {"producer_name": "tf2onnx",
                   "producer_version": __version__}
-        if self._opset > 0:
-            imp = OperatorSetIdProto()
-            imp.version = self._opset
-            kwargs["opset_imports"] = [imp]
+        opsets = []
+        imp = OperatorSetIdProto()
+        imp.version = self._opset
+        opsets.append(imp)
+        if self._extra_opset is not None:
+            opsets.extend(self._extra_opset)
+        kwargs["opset_imports"] = opsets
 
         model_proto = helper.make_model(graph, **kwargs)
 

tf2onnx/tfonnx.py

@@ -705,10 +705,24 @@ def expanddims_op7(ctx, node, name, args):
 
 
 def stridedslice_op(ctx, node, name, args):
-    # T output = StridedSlice(T input, Index begin, Index end, Index strides,
-    #               @type Index, @int begin_mask, @int end_mask, @int ellipsis_mask,
-    #               @int new_axis_mask, @int shrink_axis_mask)
-    raise ValueError("StridedSlice not implemented")
+    # only the cases strides=1 can be mapped to onnx
+    not_supported_attr = ["begin_mask", "ellipsis_mask", "end_mask", "new_axis_mask", "shrink_axis_mask"]
+    for attr_name in not_supported_attr:
+        attr = node.get_attr(attr_name)
+        if attr is not None and attr.i != 0:
+            raise ValueError("StridedSlice: attribute " + attr_name + " must be 0")
+    begin = node.inputs[1].get_tensor_value()
+    end = node.inputs[2].get_tensor_value()
+    strides = node.inputs[3].get_tensor_value()[0]
+    if strides != 1:
+        raise ValueError("StridedSlice: only strides=1 is supported")
+    node.set_attr("starts", list(begin))
+    node.set_attr("ends", list(end))
+    node.type = "Slice"
+    ctx.remove_input(node, node.input[3])
+    ctx.remove_input(node, node.input[2])
+    ctx.remove_input(node, node.input[1])
+    return node
 
 
 def pow_op(ctx, node, name, args):
@@ -765,6 +779,14 @@ def multinomial_op(ctx, node, name, args):
     return node
 
 
+def topk_op(ctx, node, name, args):
+    k = node.inputs[1].get_tensor_value()
+    node.set_attr("k", k[0])
+    node.type = "TopK"
+    ctx.remove_input(node, node.input[1])
+    return node
+
+
 # pylint: enable=W0613,C0111,W0612
 
 # map tensorflow ops to onnx ops. The format below is
@@ -847,6 +869,7 @@ def multinomial_op(ctx, node, name, args):
     "Sum": (reduce_op, ["ReduceSum"]),
     "Tanh": (direct_op, []),
     "Transpose": (transpose_op, []),
+    "TopKV2": (topk_op, []),
 }
 
 _OPSET_5 = {
@@ -1078,7 +1101,7 @@ 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):
+                     opset=None, custom_op_handlers=None, custom_rewriter=None, extra_opset=None):
     """Convert tensorflow graph to onnx graph.
         Args:
             tf_graph: tensorflow graph
@@ -1106,7 +1129,7 @@ def topological_sort(ops):
 
     onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes = tensorflow_to_onnx(tf_graph)
 
-    g = Graph(onnx_nodes, output_shapes, dtypes, target, opset)
+    g = Graph(onnx_nodes, output_shapes, dtypes, target, opset, extra_opset)
     ops = g.get_nodes()
 
     # rewrite graph

tf2onnx/utils.py

@@ -75,7 +75,7 @@
     'min', 'seed', 'ends', 'paddings', 'to', 'gamma', 'width_scale', 'normalize_variance', 'group', 'ratio', 'values',
     'dtype', 'output_shape', 'spatial', 'split', 'input_forget', 'keepdims', 'transA', 'auto_pad', 'border', 'low',
     'linear_before_reset', 'height_scale', 'output_padding', 'shape', 'kernel_shape', 'epsilon', 'size', 'starts',
-    'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales'
+    'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales', 'k'
 }