Merge pull request #399 from zhijxu-MS/use_onehot_to_map_onehot

use onnx "onehot" to map tf "onehot" if possible

Author: nbcsm

tests/test_backend.py

@@ -1128,7 +1128,7 @@ def test_onehot0(self):
             _ = tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @unittest.skip("")
+    @unittest.skip("only rank 1 is currently implemented")
     def test_onehot1(self):
         # only rank 1 is currently implemented
         x_val = np.array([[0, 2], [1, -1]], dtype=np.int32)
@@ -1139,12 +1139,40 @@ def test_onehot1(self):
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
     def test_onehot2(self):
-        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)
-        _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        for axis in [-1, 0, 1]:
+            tf.reset_default_graph()
+            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=axis, off_value=1.0, dtype=tf.float32)
+            _ = tf.identity(x_, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(9, "onehot")
+    def test_onehot3(self):
+        # rank 1
+        for np_dtype, tf_dtype in zip([np.int32, np.int64], [tf.int32, tf.int64]):
+            tf.reset_default_graph()
+            x_val = np.array([0, 1, 2, 1, 2, 0, 1, 2, 1, 2], dtype=np_dtype)
+            depth = np.array(20).astype(np.int64)
+            x = tf.placeholder(tf_dtype, x_val.shape, name=_TFINPUT)
+            on_off = np.array([5.6, 1.2]).astype(np_dtype)
+            x_ = tf.one_hot(x, depth, on_value=on_off[0], axis=-1, off_value=on_off[1])
+            _ = tf.identity(x_, name=_TFOUTPUT)
+            graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
+            self.assertTrue(len(group_nodes_by_type(graph)["OneHot"]) == 1, "onnx onehot should be used")
+        # rank 2
+        for aixs in [-1, 0, 1, 2]:
+            for np_dtype, tf_dtype in zip([np.int32, np.int64], [tf.int32, tf.int64]):
+                tf.reset_default_graph()
+                x_val = np.arange(0, 50, dtype=np_dtype).reshape([-1, 10])
+                depth = np.array(20).astype(np.int64)
+                x = tf.placeholder(tf_dtype, x_val.shape, name=_TFINPUT)
+                on_off = np.array([5.6, 1.2]).astype(np_dtype)
+                x_ = tf.one_hot(x, depth, on_value=on_off[0], axis=aixs, off_value=on_off[1])
+                _ = tf.identity(x_, name=_TFOUTPUT)
+                graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
+                self.assertTrue(len(group_nodes_by_type(graph)["OneHot"]) == 1, "onnx onehot should be used")
 
     @skip_caffe2_backend("issue undefined dim 1")
     def test_flatten0(self):

tf2onnx/tfonnx.py

@@ -1287,6 +1287,48 @@ def onehot_op(ctx, node, name, args):
         return
 
 
+def onehot_op9(ctx, node, name, args):
+    # T output = OneHot(uint8/int32/int64 input, T depth, T on-value, T off-value, @int axis, @dtype)
+    # tf requires that dtype is same as on-value's and off-value's dtype
+    # in ONNX, op's schema is (input, depth, value, @int axis), meaning of "value" is [off-value, on-value]
+    # onnxruntime only supports int64
+    output_dtype = ctx.get_dtype(node.input[2])
+    if ctx.is_target(TARGET_RS6) and output_dtype not in [onnx_pb.TensorProto.INT64, onnx_pb.TensorProto.INT32]:
+        log.warning("unsupported dtype in onnxruntime, onehot-9 can't be used directly")
+        onehot_op(ctx, node, name, args)
+        return
+
+    depth = node.input[1]
+    depth = ctx.make_node("Unsqueeze", [depth], attr={"axes": [0]}).output[0]
+
+    on_value = node.input[2]
+    off_value = node.input[3]
+    on_value = ctx.make_node("Unsqueeze", [on_value], attr={"axes": [0]}).output[0]
+    off_value = ctx.make_node("Unsqueeze", [off_value], attr={"axes": [0]}).output[0]
+    off_on_value = ctx.make_node("Concat", [off_value, on_value], attr={"axis": 0}).output[0]
+
+    indices = node.input[0]
+    if ctx.is_target(TARGET_RS6) and ctx.get_dtype(indices) != onnx_pb.TensorProto.INT64:
+        indices = ctx.make_node("Cast", [indices], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
+    node.input[0] = indices
+
+    if ctx.is_target(TARGET_RS6) and ctx.get_dtype(depth) != onnx_pb.TensorProto.INT64:
+        depth = ctx.make_node("Cast", [depth], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
+    node.input[1] = depth
+
+    if ctx.is_target(TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
+        off_on_value = ctx.make_node("Cast", [off_on_value], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
+    node.input[2] = off_on_value
+
+    del node.input[3]
+
+    if ctx.is_target(TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
+        new_node_name = utils.make_name("onehot_output")
+        new_node = ctx.insert_new_node_on_output("Cast", node.output[0], new_node_name, to=output_dtype)
+        ctx.set_dtype(new_node.output[0], output_dtype)
+        ctx.set_shape(new_node.output[0], ctx.get_shape(node.output[0]))
+
+
 def fused_batchnorm_op7(ctx, node, name, args):
     node.type = "BatchNormalization"
     # tf inputs: x, scale, bias, mean, variance
@@ -1866,6 +1908,7 @@ def where_op(ctx, node, name, args):
     "Erf": (direct_op, []),
     "Fill": (fill_op, []),
     "IsNan": (direct_op, ["IsNaN"]),
+    "OneHot": (onehot_op9, []),
     "ResizeBilinear": (upsample_op9, ["Upsample", "linear"]),
     "ResizeNearestNeighbor": (upsample_op9, ["Upsample", "nearest"]),
     "ReverseSequence": (reverse_op9, []),