Merge pull request #305 from jiafatom/customized_arg

Pass args to custom_op converter

Author: guschmue

README.md

@@ -227,7 +227,6 @@ def print_handler(ctx, node, name, args):
     #   T output = Print(T input, data, @list(type) U, @string message, @int first_n, @int summarize)
     # becomes:
     #   T output = Identity(T Input)
-    node.type = "Identity"
     node.domain = _TENSORFLOW_DOMAIN
     del node.input[1:]
     return node
@@ -239,7 +238,7 @@ with tf.Session() as sess:
     x_ = tf.Print(x, [x], "hello")
     _ = tf.identity(x_, name="output")
     onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph,
-                                                 custom_op_handlers={"Print": print_handler},
+                                                 custom_op_handlers={"Print": (print_handler, ["Identity", "mode"])},
                                                  extra_opset=[helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)],
                                                  input_names=["input:0"],
                                                  output_names=["output:0"])

tests/test_graph.py

@@ -333,8 +333,9 @@ def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
             #   T output = Print(T input, data, @list(type) U, @string message, @int first_n, @int summarize)
             # becomes:
             #   T output = Identity(T Input)
-            node.type = "Identity"
+            self.assertEqual(node.type, "Identity")
             node.domain = _TENSORFLOW_DOMAIN
+            self.assertEqual(args[0], "mode")
             del node.input[1:]
             return node
 
@@ -343,7 +344,7 @@ def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
             x_ = tf.Print(x, [x], "hello")
             _ = tf.identity(x_, name="output")
             g = process_tf_graph(sess.graph,
-                                 custom_op_handlers={"Print": print_handler},
+                                 custom_op_handlers={"Print": (print_handler, ["Identity", "mode"])},
                                  extra_opset=helper.make_opsetid(_TENSORFLOW_DOMAIN, 1))
             self.assertEqual(
                 'digraph { input1 [op_type=Placeholder shape="[2, 3]"] Print [op_type=Identity] '

tf2onnx/tfonnx.py

@@ -2284,7 +2284,7 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
     # apply custom ops on top of the assembled opset. We can either completment the opset
     # or override existing ops with a custom op.
     if custom_op_handlers is not None:
-        custom_opset = {k: [v, []] for k, v in custom_op_handlers.items()}
+        custom_opset = {k: v for k, v in custom_op_handlers.items()}
         ops_mapping.update(custom_opset)
 
     ops = g.get_nodes()