Convert lookup op

Author: RandyShuai

tf2onnx/custom_opsets/__init__.py

@@ -3,3 +3,4 @@
 """ custom tf2onnx mapping functions. """
 
 from . import ms
+from . import onnx_ml

tf2onnx/custom_opsets/onnx_ml.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+""" tf2onnx mapping functions for onnx ml domain. """
+import numpy as np
+
+from onnx import onnx_pb
+from onnx.onnx_pb import TensorProto
+from tf2onnx import constants, utils
+from tf2onnx.handler import tf_op
+
+@tf_op("HashTableV2")
+class HashTable:
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        pass
+
+@tf_op("LookupTableFindV2")
+class LookupTableFind:
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        table_node = node.inputs[0]
+        file_path = table_node.get_attr_value("shared_name")[11:-6]
+        cats_int64s = []
+        cats_strings = []
+        with open(file_path, 'r') as f:
+            for i,s in enumerate(f.readlines()):
+                cats_int64s.append(i)
+                cats_strings.append(s.strip())
+        node_inputs = node.input
+        node_outputs = node.output
+        ctx.remove_node(node.name)
+        ctx.make_node("CategoryMapper", domain=constants.AI_ONNX_ML_DOMAIN, inputs=node_inputs[1:2], outputs=node_outputs, attr={'cats_int64s':cats_int64s,'cats_strings':cats_strings})
+        customer_nodes = ctx.find_output_consumers(table_node.output[0])
+        if len(customer_nodes) == 0:
+            ctx.remove_node(table_node.name)

tf2onnx/onnx_opset/controlflow.py

@@ -365,3 +365,20 @@ def version_9(cls, ctx, node, **kwargs):
                                                        node.output[0], name=utils.make_name("where_op_added"))
         ctx.copy_shape(node.output[0], transpose_node.output[0])
         ctx.copy_dtype(node.output[0], transpose_node.output[0])
+
+@tf_op("IteratorV2")
+class Iterator:
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        ctx.remove_node(node.name)
+
+@tf_op("IteratorGetNext")
+class IteratorGetNext:
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        output_names = node.output
+        ctx.remove_node(node.name)
+        output_types = list(node.get_attr('output_types').ints)
+        output_shapes = list(node.get_attr('output_shapes').ints)
+        ctx.add_graph_input(output_names[0], output_types[0], output_shapes[:output_shapes.index(0)])
+        ctx.add_graph_input(output_names[1], output_types[1], output_shapes[output_shapes.index(0)+1:-1])

tf2onnx/tfonnx.py

@@ -99,6 +99,12 @@ def tflist_to_onnx(node_list, shape_override):
                 continue
             elif a in ignored_attr:
                 continue
+            elif a in ["key_dtype", "value_dtype", "Tin", "Tout"]:
+                attr[a] = utils.map_tf_dtype(utils.get_tf_node_attr(node, a))
+            elif a == "output_types":
+                attr[a] = [utils.map_tf_dtype(v) for v in utils.get_tf_node_attr(node, a)]
+            elif a == "output_shapes":
+                attr[a] = utils.get_tf_output_shapes_attr(node)
             else:
                 attr[a] = utils.get_tf_node_attr(node, a)
 

tf2onnx/utils.py

@@ -177,6 +177,17 @@ def get_tf_shape_attr(node):
         pass
     return dims
 
+def get_tf_output_shapes_attr(node):
+    """Get output shapes from tensorflow attr "output_shapes"."""
+    dims = [] 
+    try:
+        shapes = get_tf_node_attr(node, "output_shapes")
+        for shape in shapes:
+            dims.extend([d.size for d in shape.dim])
+            dims.append(0)
+    except:  # pylint: disable=bare-except
+        pass
+    return dims
 
 def get_tf_tensor_shape(tensor):
     shape = []