Merge pull request #786 from onnx/rashuai/gnmt2

Convert GNMT

Author: guschmue

tests/backend_test_base.py

@@ -87,6 +87,7 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
         graph_def = None
         if convert_var_to_const:
             with tf.Session() as sess:
+                tf.tables_initializer().run()
                 variables_lib.global_variables_initializer().run()
                 output_name_without_port = [n.split(':')[0] for n in output_names_with_port]
                 graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def,
@@ -96,6 +97,7 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
             tf.import_graph_def(graph_def, name='')
 
         with tf.Session() as sess:
+            tf.tables_initializer().run()
             variables_lib.global_variables_initializer().run()
             output_dict = []
             for out_name in output_names_with_port:

tests/test_backend.py

@@ -10,9 +10,11 @@
 import unittest
 from itertools import product
 
+import os
 import numpy as np
 import tensorflow as tf
 
+from tensorflow.python.ops import lookup_ops
 from backend_test_base import Tf2OnnxBackendTestBase
 # pylint reports unused-wildcard-import which is false positive, __all__ is defined in common
 from common import *  # pylint: disable=wildcard-import,unused-wildcard-import
@@ -2974,6 +2976,20 @@ def test_Conv2DBackpropInput_valid(self):
                                         name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: input_sizes_val, _INPUT1: filters_val, _INPUT2: out_backprop_val})
 
+    @check_opset_min_version(8, "CategoryMapper")
+    def test_hashtable_lookup(self):
+        filnm = "vocab.tmp"
+        words = ["apple", "pear", "banana", "cherry", "grape"]
+        query = np.array(['cherry'], dtype=object)
+        with open(filnm, "w") as f:
+            for word in words:
+                f.write(word + "\n")
+        query_holder = tf.placeholder(tf.string, shape=[len(query)], name=_TFINPUT)
+        hash_table = lookup_ops.index_table_from_file(filnm)
+        lookup_results = hash_table.lookup(query_holder)
+        self._run_test_case([lookup_results.name], {_INPUT: query})
+        os.remove(filnm)
+
 
 if __name__ == '__main__':
     unittest_main()

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,41 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+""" tf2onnx mapping functions for onnx ml domain. """
+from tf2onnx import constants
+from tf2onnx.handler import tf_op
+
+
+# pylint: disable=unused-argument,missing-docstring,unnecessary-pass
+
+@tf_op("HashTableV2")
+class HashTable:
+    @classmethod
+    def version_8(cls, ctx, node, **kwargs):
+        """ HashTable will be removed """
+        pass
+
+
+@tf_op("LookupTableFindV2")
+class LookupTableFind:
+    @classmethod
+    def version_8(cls, ctx, node, **kwargs):
+        """ convert lookup to category mapper """
+        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_name = node.name
+        node_inputs = node.input
+        node_outputs = node.output
+        ctx.remove_node(node.name)
+        new_node = ctx.make_node("CategoryMapper", domain=constants.AI_ONNX_ML_DOMAIN,
+                                 name=node_name, inputs=node_inputs[1: 2], outputs=node_outputs,
+                                 attr={'cats_int64s': cats_int64s, 'cats_strings': cats_strings})
+        ctx.set_shape(new_node.name + ":0", [-1])
+        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,22 @@ 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_8(cls, ctx, node, **kwargs):
+        ctx.remove_node(node.name)
+
+@tf_op("IteratorGetNext")
+class IteratorGetNext:
+    @classmethod
+    def version_8(cls, ctx, node, **kwargs):
+        output_names = node.output
+        type_0 = ctx.get_dtype(output_names[0])
+        type_1 = ctx.get_dtype(output_names[1])
+        shape_0 = ctx.get_shape(output_names[0])
+        shape_1 = ctx.get_shape(output_names[1])
+        ctx.remove_node(node.name)
+        ctx.add_graph_input(output_names[0], type_0, shape_0)
+        ctx.add_graph_input(output_names[1], type_1, shape_1)

tf2onnx/tfonnx.py

@@ -46,7 +46,9 @@ def tflist_to_onnx(node_list, shape_override):
     ignored_attr = ["unknown_rank", "_class", "Tshape", "use_cudnn_on_gpu", "Index", "Tpaddings",
                     "TI", "Tparams", "Tindices", "Tlen", "Tdim", "dynamic_size", "Tmultiples",
                     "Tblock_shape", "Tcrops", "index_type", "Taxis", "U", "maxval",
-                    "Tout", "Tlabels", "Tindex", "element_shape", "Targmax", "T_threshold"]
+                    "Tout", "Tlabels", "Tindex", "element_shape", "Targmax", "T_threshold",
+                    "output_types", "output_shapes", "key_dtype", "value_dtype", "Tin", "Tout"]
+
     # some stats
     op_cnt = collections.Counter()
     attr_cnt = collections.Counter()