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

Author: zhijxu-MS

tests/test_backend.py

@@ -1146,6 +1146,31 @@ def test_onehot2(self):
         _ = 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 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=-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")
+
     @skip_caffe2_backend("issue undefined dim 1")
     def test_flatten0(self):
         x_val = np.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]], dtype=np.float32)

tf2onnx/tfonnx.py

@@ -1287,6 +1287,45 @@ 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, [off-value, on-value], @int axis)
+    # onnxruntime only supports int64
+    output_dtype = ctx.get_dtype(node.input[2])
+    if output_dtype not in [onnx_pb.TensorProto.INT64, onnx_pb.TensorProto.INT32]:
+        log.warning("unsupported dtype in onnxruntime, onehot can't be used directly")
+        return onehot_op(ctx, node, name, args)
+
+    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.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.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 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 output_dtype != onnx_pb.TensorProto.INT64:
+        new_output = utils.make_name("onehot_output")
+        ctx.insert_new_node_on_output("Cast", node.output[0], new_output, to=output_dtype)
+
+
 def fused_batchnorm_op7(ctx, node, name, args):
     node.type = "BatchNormalization"
     # tf inputs: x, scale, bias, mean, variance
@@ -1865,6 +1904,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, []),