Made from_keras work for tf1 (#1529)

* Made from_keras work for tf1

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

* pylint

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

Author: TomWildenhain-Microsoft

tests/test_api.py

@@ -13,6 +13,7 @@
 from onnx import helper
 
 from common import check_tf_min_version, unittest_main, requires_custom_ops, check_opset_min_version
+from tf2onnx.tf_loader import is_tf2
 from backend_test_base import Tf2OnnxBackendTestBase
 import tf2onnx
 
@@ -67,16 +68,16 @@ def _test_keras_api(self, large_model=False):
         ky1 = model.predict([x, n])
         self.assertAllClose(ky1, oy[0], rtol=0.3, atol=0.1)
 
-    @check_tf_min_version("2.0")
+    @check_tf_min_version("1.15")
     def test_keras_api(self):
         self._test_keras_api(large_model=False)
 
-    @check_tf_min_version("2.2")
+    @check_tf_min_version("1.15")
     def test_keras_api_large(self):
         self._test_keras_api(large_model=True)
 
     @requires_custom_ops()
-    @check_tf_min_version("2.0")
+    @check_tf_min_version("1.15")
     @check_opset_min_version(11, "SparseToDense")
     def test_keras_hashtable(self):
 
@@ -101,6 +102,8 @@ def test_keras_hashtable(self):
 
         inp1 = np.array([[2.], [3.]], dtype=np.float32)
         inp2 = np.array([["a"], ["b"]], dtype=np.str)
+        if not is_tf2():
+            tf.keras.backend.get_session().run(tf.tables_initializer(name='init_all_tables'))
         k_res = model.predict([inp1, inp2])
         spec = (tf.TensorSpec((None, 1), dtype=tf.float32, name="f_inp"),
                 tf.TensorSpec((None, 1), tf.string, name="s_inp"))

tf2onnx/convert.py

@@ -284,6 +284,42 @@ def tensor_names_from_structed(concrete_func, input_names, output_names):
     return tensors_to_rename
 
 
+def _from_keras_tf1(model, input_signature=None, opset=None, custom_ops=None, custom_op_handlers=None,
+                    custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None,
+                    target=None, large_model=False, output_path=None):
+    """from_keras for tf 1.15"""
+
+    input_names = [t.name for t in model.inputs]
+    output_names = [t.name for t in model.outputs]
+    tensors_to_rename = dict(zip(input_names, model.input_names))
+    if len(set(model.output_names)) == len(model.output_names):
+        # In very rare cases, keras has a bug where it will give multiple outputs the same name
+        tensors_to_rename.update(zip(output_names, model.output_names))
+
+    sess = tf.keras.backend.get_session(model.outputs)
+
+    with tf.device("/cpu:0"):
+        frozen_graph, initialized_tables = tf_loader.freeze_session(sess, input_names, output_names, get_tables=True)
+        model_proto, external_tensor_storage = _convert_common(
+            frozen_graph,
+            name=model.name,
+            continue_on_error=True,
+            target=target,
+            opset=opset,
+            custom_op_handlers=custom_ops,
+            extra_opset=extra_opset,
+            shape_override=shape_override,
+            input_names=input_names,
+            output_names=output_names,
+            inputs_as_nchw=inputs_as_nchw,
+            large_model=large_model,
+            tensors_to_rename=tensors_to_rename,
+            initialized_tables=initialized_tables,
+            output_path=output_path)
+
+        return model_proto, external_tensor_storage
+
+
 def from_keras(model, input_signature=None, opset=None, custom_ops=None, custom_op_handlers=None,
                custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None,
                target=None, large_model=False, output_path=None):
@@ -306,7 +342,8 @@ def from_keras(model, input_signature=None, opset=None, custom_ops=None, custom_
         An ONNX model_proto and an external_tensor_storage dict.
     """
     if LooseVersion(tf.__version__) < "2.0":
-        raise NotImplementedError("from_keras requires tf-2.0 or newer")
+        return _from_keras_tf1(model, input_signature, opset, custom_ops, custom_op_handlers, custom_rewriter,
+                               inputs_as_nchw, extra_opset, shape_override, target, large_model, output_path)
 
     from tensorflow.python.keras.saving import saving_utils as _saving_utils # pylint: disable=import-outside-toplevel
 

tf2onnx/tf_loader.py

@@ -214,7 +214,7 @@ def from_function(func, input_names, output_names, large_model=False):
     return graph_def
 
 
-def freeze_session(sess, input_names=None, output_names=None):
+def freeze_session(sess, input_names=None, output_names=None, get_tables=False):
     """Freezes the state of a session into a pruned computation graph."""
     output_node_names = [i.split(':')[:-1][0] for i in output_names]
     keep_var_names = [i.split(':')[:-1][0] for i in input_names]
@@ -226,6 +226,19 @@ def freeze_session(sess, input_names=None, output_names=None):
         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)
+        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)
+                try:
+                    k, v = lookup_ops.lookup_table_export_v2(h, k_dtype, val_dtype)
+                    k, v = sess.run([k, v])
+                    initialized_tables[n] = (k, v)
+                except Exception:  # pylint: disable=broad-except
+                    logger.warning("Could not initialize table with shared_name = %r", n)
+            return graph_def, initialized_tables
     return graph_def
 
 
@@ -348,18 +361,8 @@ def _from_saved_model_v1(sess, model_path, input_names, output_names, tag, signa
                 if output_tensor.name not in output_names:
                     output_names.append(output_tensor.name)
                     tensors_to_rename[output_tensor.name] = structured_name
-    frozen_graph = freeze_session(sess, input_names=input_names, output_names=output_names)
-    table_names, key_dtypes, value_dtypes = get_hash_table_info(frozen_graph)
-    initialized_tables = {}
-    tf.tables_initializer().run()
-    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)
-        try:
-            k, v = lookup_ops.lookup_table_export_v2(h, k_dtype, val_dtype)
-            k, v = sess.run([k, v])
-            initialized_tables[n] = (k, v)
-        except Exception:  # pylint: disable=broad-except
-            logger.warning("Could not initialize table with shared_name = %r", n)
+    frozen_graph, initialized_tables = \
+        freeze_session(sess, input_names=input_names, output_names=output_names, get_tables=True)
     return frozen_graph, input_names, output_names, initialized_tables, tensors_to_rename