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

- add support for onehot (rank=1 for now)
- onnx optset < 8 does define reshape for float only ... insert casts for now
- update tests
- pass backend for unittest via command line

fix for #33

Author: guschmue

tests/run_pretrained_models.py

@@ -271,8 +271,18 @@ 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())
+
         g = tf.import_graph_def(graph_def, name='')
         with tf.Session(graph=g) as sess:
+
+            # fix inputs if needed
+            for k in inputs.keys():
+                t = sess.graph.get_tensor_by_name(k)
+                dtype = tf.as_dtype(t.dtype).name
+                if type != "float32":
+                    v = inputs[k]
+                    inputs[k] = v.astype(dtype)
+
             tf_results = self.run_tensorflow(sess, inputs)
             onnx_graph = None
             print("\ttensorflow", "OK")

tests/test_backend.py

@@ -1,7 +1,9 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
 
+import argparse
 import os
+import sys
 import tempfile
 import unittest
 from collections import namedtuple
@@ -553,6 +555,16 @@ def test_reshape(self):
         self.assertEqual(expected.shape, actual.shape)
         self.assertAllClose(expected, actual)
 
+    def test_reshape_int(self):
+        x_val = np.array([1, 2, 3, 4], dtype=np.int32).reshape((2, 2))
+        x = tf.placeholder(tf.int32, [2, 2], name=_TFINPUT)
+        shape = tf.constant([1, 4])
+        x_ = tf.reshape(x, shape)
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertEqual(expected.shape, actual.shape)
+        self.assertAllClose(expected, actual)
+
     @unittest.skipIf(OPSET < 5 or BACKEND in ["onnxmsrtnext"], "since opset 5, broken in msrtnext")
     def test_reshape_dynamic(self):
         x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2))
@@ -788,13 +800,35 @@ def test_cast(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
-    @unittest.skip
-    def test_onehot(self):
+    def test_onehot0(self):
         # no such op in onnx
         x_val = np.array([0, 1, 2], dtype=np.int32)
+        depth = 5
+        for axis in [-1, 0, 1]:
+            tf.reset_default_graph()
+            x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
+            x_ = tf.one_hot(x, depth, on_value=5.0, axis=axis, off_value=1.0, dtype=tf.float32)
+            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_onehot1(self):
+        # no such op in onnx
+        x_val = np.array([[0, 2], [1, -1]], dtype=np.int32)
         depth = 3
         x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
-        x_ = tf.one_hot(x, depth, on_value=1, axis=0, off_value=0)
+        x_ = tf.one_hot(x, depth, on_value=5.0, axis=-1, off_value=0.0, dtype=tf.float32)
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
+    def test_onehot2(self):
+        # no such op in onnx
+        x_val = np.array([0, 1, 2, 1, 2, 0, 1, 2, 1, 2], dtype=np.int32)
+        depth = 20
+        x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
+        x_ = tf.one_hot(x, depth, on_value=5.0, axis=-1, off_value=1.0, dtype=tf.float32)
         output = tf.identity(x_, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
@@ -854,14 +888,6 @@ def test_unstack_axis(self):
             actual, expected = self._run(output, {}, {})
             self.assertAllClose(expected, actual)
 
-    def test_unstack_axis1(self):
-        x_val = np.random.randn(10, 3, 4).astype("float32")
-        x = tf.constant(x_val, dtype=tf.float32)
-        x_ = tf.unstack(x, axis=1)
-        output = tf.identity(x_, name=_TFOUTPUT)
-        actual, expected = self._run(output, {}, {})
-        self.assertAllClose(expected, actual)
-
     @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "Space2Depth not implemented, works on onnxmsrtnext")
     def test_space_to_depth(self):
         x_val = make_xval([1, 2, 2, 1])
@@ -871,7 +897,6 @@ def test_space_to_depth(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
-
     @unittest.skipIf(OPSET < 6, "supported since opset 6")
     def test_addn(self):
         x_val = np.arange(3*2*3).astype("float32")
@@ -935,4 +960,14 @@ def test_fill(self):
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--backend', default='caffe2',
+                        choices=["caffe2", "onnxmsrt", "onnxmsrtnext", "onnx-tensorflow"],
+                        help="backend to test against")
+    parser.add_argument('unittest_args', nargs='*')
+
+    args = parser.parse_args()
+    BACKEND = args.backend
+    # Now set the sys.argv to the unittest_args (leaving sys.argv[0] alone)
+    sys.argv[1:] = args.unittest_args
     unittest.main()

tests/unity.yaml

@@ -42,13 +42,11 @@ BananaRL:
     - value_estimate:0
 
 Basic:
-  # needs: onehot
-  disabled: true
   url: https://github.com/Unity-Technologies/ml-agents/raw/master/unity-environment/Assets/ML-Agents/Examples/Basic/TFModels/Basic.bytes
   model: Basic.bytes
   input_get: get_random
   inputs:
-    "vector_observation:0": [1, 1]
+    "vector_observation:0": [10, 1]
   outputs:
     - action:0
     - action_probs:0

tf2onnx/graph.py

@@ -129,6 +129,15 @@ def shape(self):
                 shape[0] = utils.ONNX_UNKNOWN_DIMENSION
         return shape
 
+    def get_tensor_type(self):
+        """Get the onnx data type of a tensor."""
+        t = self.get_attr("value")
+        if t:
+            t = helper.get_attribute_value(t)
+            if t:
+                return utils.ONNX_TO_NUMPY_DTYPE[t.data_type]
+        return onnx_pb.TensorProto.FLOAT
+
     def get_tensor_value(self):
         """Get value for onnx tensor."""
         if not self.is_const():
@@ -478,7 +487,8 @@ def insert_new_node_on_input(self, node, op_type, input_name, name=None, **kwarg
         Returns:
             node that was inserted
         """
-        assert isinstance(input_name, str) and isinstance(op_type, str)
+        if name is None:
+            name = utils.make_name(node.name)
         new_output = name + ":0"
         new_node = Node(helper.make_node(op_type, [input_name], [new_output], name=name, **kwargs), self)
         for i, n in enumerate(node.input):
@@ -504,6 +514,14 @@ def insert_new_node_on_output(self, op_type, output_name, name=None, **kwargs):
         self.replace_all_inputs(self.get_nodes(), output_name, new_output)
         return new_node
 
+    def find_output_consumers(self, output_name):
+        """Find all nodes consuming a given output."""
+        nodes = []
+        for node in self.get_nodes():
+            if output_name in node.input:
+                nodes.append(node)
+        return nodes
+
     @staticmethod
     def replace_all_inputs(ops, old_input, new_input):
         """Replace all inputs pointing to old_input with new_input."""

tf2onnx/tfonnx.py

@@ -29,6 +29,7 @@
 POSSIBLE_TARGETS = [TARGET_RS4, TARGET_CAFFE2]
 DEFAULT_TARGET = [TARGET_RS4, TARGET_CAFFE2]
 
+
 def tensorflow_to_onnx(graph):
     """
     Load tensorflow graph into an onnx graph with minimal rewrites so
@@ -119,10 +120,9 @@ def _convert_shapenode_to_int64(ctx, node, input_number):
         shape_node.set_attr("value", onnx_tensor)
         return [node]
     else:
-        op_name = utils.make_name(node.name)
-        cast_op = ctx.insert_new_node_on_input(node, "Cast", name, name=op_name)
+        cast_op = ctx.insert_new_node_on_input(node, "Cast", name)
         cast_op.set_attr("to", onnx_pb.TensorProto.INT64)
-        ctx.copy_shape(name, op_name + ":0")
+        ctx.copy_shape(name, cast_op.output[0])
         return [cast_op, node]
 
 # pylint: disable=W0613,C0111,W0612
@@ -274,8 +274,29 @@ def reshape_op(ctx, node, name, args):
 
 
 def reshape_op5(ctx, node, name, args):
+    need_casting = node.dtype in [onnx_pb.TensorProto.INT32,
+                                  onnx_pb.TensorProto.INT16,
+                                  onnx_pb.TensorProto.INT64]
     # onnx wants reshape.input[1] to have the value be int64 which is not the case for tensorflow.
-    return _convert_shapenode_to_int64(ctx, node, 1)
+    nodes = _convert_shapenode_to_int64(ctx, node, 1)
+    if not need_casting:
+        # onnx reshape can handle the type - done
+        return nodes
+
+    # onnx < opset 8 does not know reshape for other types than float*, wrap the reshape in casts
+    input_cast = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
+    input_cast.set_attr("to", onnx_pb.TensorProto.FLOAT)
+    ctx.copy_shape(name, input_cast.output[0])
+
+    # if the next node is already a cast we don't need to insert another one
+    next_nodes = ctx.find_output_consumers(node.output[0])
+    if len(next_nodes) != 1 or next_nodes[0].type != "Cast":
+        op_name = utils.make_name(node.name)
+        output_cast = ctx.insert_new_node_on_output("Cast", node.output[0], name=op_name)
+        output_cast.set_attr("to", node.dtype)
+        ctx.copy_shape(name, output_cast.output[0])
+        nodes.append(output_cast)
+    return [input_cast] + nodes
 
 
 NCHW_TO_NHWC = [0, 2, 3, 1]
@@ -317,8 +338,7 @@ def calc_shape(a, b):
         else:
             # if input comes from a op, insert transpose op
             input_name = node.input[0]
-            op_name = utils.make_name(node.name)
-            transpose = ctx.insert_new_node_on_input(node, "Transpose", input_name, name=op_name)
+            transpose = ctx.insert_new_node_on_input(node, "Transpose", input_name)
             transpose.set_attr("perm", NHWC_TO_NCHW)
             transpose.inserted_nchw = True
             ctx.set_shape(transpose.output[0], calc_shape(ctx.get_shape(input_name), NHWC_TO_NCHW))
@@ -336,9 +356,8 @@ def calc_shape(a, b):
                 parent.data_format = "NCHW"
         else:
             # kernel comes from op, insert transpose op
-            op_name = utils.make_name(node.name)
             input_name = node.input[1]
-            transpose = ctx.insert_new_node_on_input(node, "Transpose", input_name, name=op_name)
+            transpose = ctx.insert_new_node_on_input(node, "Transpose", input_name)
             transpose.set_attr("perm", HWCN_TO_NCHW)
             transpose.inserted_nchw = True
             ctx.copy_shape(input_name, transpose.output[0])
@@ -349,18 +368,16 @@ def calc_shape(a, b):
         if new_kernel_shape:
             if ctx.opset < 5:
                 # old reshape takes new shape as attribute
-                op_name = utils.make_name(node.name)
                 input_name = node.input[1]
-                reshape = ctx.insert_new_node_on_input(node, "Reshape", input_name, name=op_name)
+                reshape = ctx.insert_new_node_on_input(node, "Reshape", input_name)
                 reshape.set_attr("shape", new_kernel_shape)
                 ctx.set_shape(reshape.output[0], new_kernel_shape)
             else:
                 # new reshape takes new shape as input[1]
-                op_name = utils.make_name(node.name)
                 shape_name = utils.make_name(node.name)
                 shape_node = ctx.make_const(shape_name, "Const", np.array(new_kernel_shape, dtype=np.int64))
                 input_name = node.input[1]
-                reshape = ctx.insert_new_node_on_input(node, "Reshape", input_name, name=op_name)
+                reshape = ctx.insert_new_node_on_input(node, "Reshape", input_name)
                 reshape.input.append(shape_name)
                 ctx.set_shape(reshape.output[0], new_kernel_shape)
             nodes.append(reshape)
@@ -820,7 +837,7 @@ def minmax_op(ctx, node, name, args):
             input_node = node.inputs[i]
             dtype = ctx.dtypes[node.input[i]]
             zero_name = utils.make_name(input_node.name)
-            zero_node = ctx.make_const(zero_name, "Const", np.zeros(shapeo, dtype=utils.ONNX_TO_NUMPY_DTYPE[dtype]))
+            ctx.make_const(zero_name, "Const", np.zeros(shapeo, dtype=utils.ONNX_TO_NUMPY_DTYPE[dtype]))
             op_name = utils.make_name(input_node.name)
             output_name = op_name + ":0"
             add_node = Node(helper.make_node("Add", [input_node.output[0], zero_name],
@@ -853,6 +870,7 @@ def pack_op(ctx, node, name, args):
     ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], output_name)
     return [concat] + nodes
 
+
 def unpack_op(ctx, node, name, args):
     # hack to make up for the missing onnx unpack op
     axis = node.get_attr("axis").i
@@ -870,6 +888,44 @@ def unpack_op(ctx, node, name, args):
     return nodes
 
 
+def onehot_op(ctx, node, name, args):
+    # until there is no onehot op in onnx, a workaround using gather from eye
+    indices_name = node.input[0]
+    indices_shape = ctx.get_shape(indices_name)
+    if len(indices_shape) != 1:
+        # TODO: this works for rank=1 but tensorflow supports more than this.
+        # Same principle should work but we need to implemtn our own eye.
+        raise ValueError("onehot op: only rank1 is supported")
+    axis = node.get_attr("axis")
+    # axis becomes axis for gather
+    node.set_attr("axis", 0)
+    depth = node.inputs[1].get_tensor_value()[0]
+    on = node.inputs[2].get_tensor_value()[0]
+    off = node.inputs[3].get_tensor_value()[0]
+    dtype = node.inputs[2].get_tensor_type()
+    eye = np.eye(depth, dtype=dtype)
+    if on != 0:
+        eye[eye == 1] = on
+        eye[eye == 0] = off
+    else:
+        eye[eye == 0] = off
+        eye[eye == 1] = on
+    const_name = utils.make_name(node.name)
+    ctx.make_const(const_name, "Const", eye)
+    # setup gather inputs
+    del node.input[:]
+    node.input.append(const_name)
+    node.input.append(indices_name)
+    node.type = "Gather"
+    if axis.i == 0:
+        # TODO: revisit for rank > 1
+        name = utils.make_name(node.name)
+        transpose_op = ctx.insert_new_node_on_output("Transpose", node.output[0], name)
+        ctx.copy_shape(node.output[0], transpose_op.output[0])
+        return [node, transpose_op]
+    return node
+
+
 # pylint: enable=W0613,C0111,W0612
 
 # map tensorflow ops to onnx ops. The format below is
@@ -962,6 +1018,7 @@ def unpack_op(ctx, node, name, args):
 _OPSET_5 = {
     "Reshape": (reshape_op5, []),
     "ExpandDims": (expanddims_op7, []),
+    "OneHot": (onehot_op, []),
 }
 
 _OPSET_6 = {
@@ -1183,7 +1240,7 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
 def tf_optimize(sess, inputs, outputs, graph_def):
     """Optimize tensorflow graph for inference."""
     transforms = [
-        #"fold_constants(ignore_errors=true)",
+        "fold_constants(ignore_errors=true)",
         "fold_batch_norms",
         "fold_old_batch_norms",
     ]