Improve lookup table loading for from_keras conversion

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/backend_test_base.py

@@ -183,11 +183,11 @@ def freeze_and_run_tf(self, func, feed_dict, outputs, as_session, premade_placeh
                 graph_def = freeze_session(sess,
                                            input_names=list(feed_dict.keys()),
                                            output_names=outputs)
-                table_names, key_dtypes, value_dtypes = get_hash_table_info(graph_def)
+                table_info = get_hash_table_info(graph_def)
                 initialized_tables = {}
-                for n, k_dtype, val_dtype in zip(table_names, key_dtypes, value_dtypes):
-                    h = lookup_ops.hash_table_v2(k_dtype, val_dtype, shared_name=n)
-                    k, v = lookup_ops.lookup_table_export_v2(h, k_dtype, val_dtype)
+                for info in table_info:
+                    h = lookup_ops.hash_table_v2(info.key_dtype, info.val_dtype, shared_name=info.shared_name)
+                    k, v = lookup_ops.lookup_table_export_v2(h, info.key_dtype, info.val_dtype)
                     initialized_tables[n] = (sess.run(k), sess.run(v))
 
             tf_reset_default_graph()

tf2onnx/tf_loader.py

@@ -4,6 +4,7 @@
 """Methods to load tensorflow graph from graphdef, checkpoint or saved_model."""
 
 import logging
+import uuid
 from distutils.version import LooseVersion
 
 import tensorflow as tf
@@ -15,7 +16,8 @@
 from tensorflow.python.util import compat
 
 from tf2onnx import utils
-from tf2onnx.tf_utils import get_tf_version, tflist_to_onnx, get_hash_table_info, replace_placeholders_with_tables
+from tf2onnx.tf_utils import (get_tf_version, tflist_to_onnx, get_hash_table_info, replace_placeholders_with_tables,
+                              HashTableInfo)
 
 logger = logging.getLogger(__name__)
 
@@ -184,7 +186,7 @@ def from_trackable(trackable, concrete_func, inputs, outputs, large_model):
     err_large_model = "model exceeds maximum protobuf size of 2GB. Try setting large_model."
 
     # Avoid errors due to bug in TF freezing
-    removed_resource_to_placeholder, graph_captures_copy, func_captures_copy = \
+    removed_resource_to_placeholder, placeholder_to_resource, graph_captures_copy, func_captures_copy = \
         _remove_non_variable_resources_from_captures(concrete_func)
 
     try:
@@ -197,16 +199,28 @@ def from_trackable(trackable, concrete_func, inputs, outputs, large_model):
     # We might be returning the concrete_func so let's put it back in working order
     _restore_captured_resources(concrete_func, graph_captures_copy, func_captures_copy)
 
-    table_names, key_dtypes, value_dtypes = get_hash_table_info(frozen_graph)
+    table_info = get_hash_table_info(frozen_graph)
     placeholder_to_table_info = {}
-    _get_hash_table_info_from_trackable(trackable, table_names, key_dtypes, value_dtypes,
+    _get_hash_table_info_from_trackable(trackable, table_info,
                                         removed_resource_to_placeholder, placeholder_to_table_info)
 
     initialized_tables = {}
-    for n, k_dtype, val_dtype in zip(table_names, key_dtypes, value_dtypes):
-        h = lookup_ops.hash_table_v2(k_dtype, val_dtype, shared_name=n)
+    for info in table_info:
+        if info.shared_name is not None:
+            h = lookup_ops.hash_table_v2(info.key_dtype, info.val_dtype, shared_name=info.shared_name)
+            n = info.shared_name
+        elif info.resource_input in placeholder_to_resource and info.resource_input not in placeholder_to_table_info:
+            # We found a lookup op with no corresponding HashTable op, but we can associate the placeholder input
+            # from the op with the resource handle from graph captures and make up a shared_name
+            h = placeholder_to_resource[info.resource_input]
+            n = str(uuid.uuid4()).encode()
+            info.shared_name = n
+            placeholder_to_table_info[info.resource_input] = info
+        else:
+            # Found a lookup op but the corresponding HashTable op has already been found and processed.
+            continue
         try:
-            k, v = lookup_ops.lookup_table_export_v2(h, k_dtype, val_dtype)
+            k, v = lookup_ops.lookup_table_export_v2(h, info.key_dtype, info.val_dtype)
             initialized_tables[n] = (k.numpy(), v.numpy())
         except Exception:  # pylint: disable=broad-except
             logger.warning("Could not initialize table with shared_name = %r", n)
@@ -260,14 +274,14 @@ def freeze_session(sess, input_names=None, output_names=None, get_tables=False):
         for node in graph_def.node:
             node.device = ""
         graph_def = convert_variables_to_constants(sess, graph_def, output_node_names)
-        table_names, key_dtypes, value_dtypes = get_hash_table_info(graph_def)
+        table_info = get_hash_table_info(graph_def)
         if get_tables:
             initialized_tables = {}
             tf.tables_initializer().run(session=sess)
-            for n, k_dtype, val_dtype in zip(table_names, key_dtypes, value_dtypes):
-                h = lookup_ops.hash_table_v2(k_dtype, val_dtype, shared_name=n)
+            for info in table_info:
+                h = lookup_ops.hash_table_v2(info.key_dtype, info.val_dtype, shared_name=info.shared_name)
                 try:
-                    k, v = lookup_ops.lookup_table_export_v2(h, k_dtype, val_dtype)
+                    k, v = lookup_ops.lookup_table_export_v2(h, info.key_dtype, info.val_dtype)
                     k, v = sess.run([k, v])
                     initialized_tables[n] = (k, v)
                 except Exception:  # pylint: disable=broad-except
@@ -403,7 +417,7 @@ def _from_saved_model_v1(sess, model_path, input_names, output_names, tag, signa
     return frozen_graph, input_names, output_names, initialized_tables, tensors_to_rename
 
 
-def _get_hash_table_info_from_trackable(trackable, table_names, key_dtypes, value_dtypes,
+def _get_hash_table_info_from_trackable(trackable, table_info,
                                         removed_resource_to_placeholder, placeholder_to_table_info):
     # pylint: disable=protected-access
     stack = [trackable]
@@ -420,26 +434,22 @@ def _get_hash_table_info_from_trackable(trackable, table_names, key_dtypes, valu
             continue
         for t in r.__dict__.values() if hasattr(r, '__dict__') else []:
             if isinstance(t, TfStaticHashTableType) and hasattr(t, '_shared_name'):
-                table_names.append(t._shared_name.encode())
-                key_dtypes.append(t.key_dtype.as_datatype_enum)
-                value_dtypes.append(t.value_dtype.as_datatype_enum)
+                info = HashTableInfo(t._shared_name.encode(), t.key_dtype.as_datatype_enum,
+                                     t.value_dtype.as_datatype_enum)
+                table_info.append(info)
                 table_handle = id(t.resource_handle)
                 if table_handle in removed_resource_to_placeholder:
-                    table_info = (table_names[-1], key_dtypes[-1], value_dtypes[-1])
-                    placeholder_to_table_info[removed_resource_to_placeholder[table_handle]] = table_info
+                    placeholder_to_table_info[removed_resource_to_placeholder[table_handle]] = info
         if isinstance(r, TfRestoredResourceType) and hasattr(r, '_create_resource'):
             try:
                 table_handle = id(r.resource_handle)
             except Exception:  # pylint: disable=broad-except
                 continue
             initializer = r._create_resource.concrete_functions[0].function_def
-            new_names, new_k_dtypes, new_v_dtypes = get_hash_table_info(initializer.node_def)
-            table_names.extend(new_names)
-            key_dtypes.extend(new_k_dtypes)
-            value_dtypes.extend(new_v_dtypes)
-            if table_handle in removed_resource_to_placeholder and len(new_names) == 1:
-                table_info = (new_names[0], new_k_dtypes[0], new_v_dtypes[0])
-                placeholder_to_table_info[removed_resource_to_placeholder[table_handle]] = table_info
+            new_table_info = get_hash_table_info(initializer.node_def)
+            table_info.extend(new_table_info)
+            if table_handle in removed_resource_to_placeholder and len(new_table_info) == 1:
+                placeholder_to_table_info[removed_resource_to_placeholder[table_handle]] = new_table_info[0]
 
 
 def _remove_non_variable_resources_from_captures(concrete_func):
@@ -449,6 +459,7 @@ def _remove_non_variable_resources_from_captures(concrete_func):
     """
     # pylint: disable=protected-access
     resource_id_to_placeholder = {}
+    placeholder_to_resource = {}
     graph_captures_copy = None
     func_captures_copy = None
     if hasattr(concrete_func.graph, '_captures') and hasattr(concrete_func, '_captured_inputs'):
@@ -459,6 +470,7 @@ def _remove_non_variable_resources_from_captures(concrete_func):
             val_tensor, name_tensor = v
             if val_tensor.dtype == tf.resource and id(val_tensor) not in variable_handles:
                 resource_id_to_placeholder[id(val_tensor)] = name_tensor.name.split(':')[0]
+                placeholder_to_resource[name_tensor.name.split(':')[0]] = val_tensor
                 del concrete_func.graph._captures[k]
                 for i in reversed(range(len(concrete_func._captured_inputs))):
                     if concrete_func._captured_inputs[i] is val_tensor:
@@ -472,7 +484,7 @@ def _remove_non_variable_resources_from_captures(concrete_func):
     else:
         logger.warning(
             "Could not search for non-variable resources. Concrete function internal representation may have changed.")
-    return resource_id_to_placeholder, graph_captures_copy, func_captures_copy
+    return resource_id_to_placeholder, placeholder_to_resource, graph_captures_copy, func_captures_copy
 
 
 def _restore_captured_resources(concrete_func, graph_captures_copy, func_captures_copy):

tf2onnx/tf_utils.py

@@ -278,6 +278,13 @@ def is_huge_shape(x):
     logger.info("Computed %d values for constant folding", len(outputs_to_values))
     return outputs_to_values, outputs_to_dtypes
 
+class HashTableInfo:
+    def __init__(self, shared_name, key_dtype, val_dtype, resource_input=None):
+        self.shared_name = shared_name
+        self.key_dtype = key_dtype
+        self.val_dtype = val_dtype
+        self.resource_input = resource_input
+
 def get_hash_table_info(nodes_or_graph_def):
     """
     Return lists of the shared_names, key_dtypes, and value_dtypes of all hash tables declared in the graph_def
@@ -287,18 +294,16 @@ def get_hash_table_info(nodes_or_graph_def):
         nodes = nodes_or_graph_def.node
     else:
         nodes = nodes_or_graph_def
-    names = []
-    key_dtypes = []
-    val_dtypes = []
+    info = []
     for n in nodes:
+        if n.op == "LookupTableFindV2":
+            info.append(HashTableInfo(None, n.attr['Tin'].type, n.attr['Tout'].type, n.input[0]))
         if n.op in ["HashTableV2", "MutableHashTableV2"]:
             if all(k in n.attr for k in ['shared_name', 'key_dtype', 'value_dtype']):
                 name = n.attr['shared_name'].s
                 if name != b'':
-                    names.append(name)
-                    key_dtypes.append(n.attr['key_dtype'].type)
-                    val_dtypes.append(n.attr['value_dtype'].type)
-    return names, key_dtypes, val_dtypes
+                    info.append(HashTableInfo(name, n.attr['key_dtype'].type, n.attr['value_dtype'].type))
+    return info
 
 def replace_placeholders_with_tables(graph_def, placeholder_to_table_info):
     """
@@ -307,13 +312,13 @@ def replace_placeholders_with_tables(graph_def, placeholder_to_table_info):
     """
     for n in graph_def.node:
         if n.op == "Placeholder" and n.name in placeholder_to_table_info:
-            name, key_dtype, val_dtype = placeholder_to_table_info[n.name]
+            info = placeholder_to_table_info[n.name]
             for a in list(n.attr):
                 del n.attr[a]
             n.op = "HashTableV2"
-            n.attr['shared_name'].s = name
-            n.attr['key_dtype'].type = key_dtype
-            n.attr['value_dtype'].type = val_dtype
+            n.attr['shared_name'].s = info.shared_name
+            n.attr['key_dtype'].type = info.key_dtype
+            n.attr['value_dtype'].type = info.val_dtype
 
 def read_tf_node_def_attrs(node_def, input_dtypes, input_shapes):
     """Given a tf node def, returns a dict of attribute names to values"""