Merge branch 'master' of https://github.com/onnx/tensorflow-onnx into profile

Author: sdpython

README.md

@@ -193,6 +193,12 @@ Only valid with parameter `--saved_model`. Specifies which signature to use with
 
 Only valid with parameter `--saved_model`. If a model contains a list of concrete functions, under the function name `__call__` (as can be viewed using the command `saved_model_cli show --all`), this parameter is a 0-based integer specifying which function in that list should be converted. This parameter takes priority over `--signature_def`, which will be ignored.
 
+#### --large_model
+
+(This is experimental, valid only for TF2.x models)
+
+Only valid with parameter `--saved_model`. When set, creates a zip file containing the ONNX protobuf model and large tensor values stored externally. This allows for converting models that exceed the 2 GB protobuf limit.
+
 #### --target
 
 Some models require special handling to run on some runtimes. In particular, the model may use unsupported data types. Workarounds are activated with ```--target TARGET```. Currently supported values are listed on this [wiki](https://github.com/onnx/tensorflow-onnx/wiki/target). If your model will be run on Windows ML, you should specify the appropriate target value.
@@ -274,7 +280,8 @@ tf2onnx.tfonnx.process_tf_graph(tf_graph,
             opset=None, custom_op_handlers=None,
             custom_rewriter=None, extra_opset=None,
             shape_override=None, inputs_as_nchw=None,
-            input_names=None, output_names=None):
+            input_names=None, output_names=None,
+            const_node_values=None):
     """Convert tensorflow graph to onnx graph.
         Args:
             tf_graph: tensorflow graph
@@ -289,6 +296,7 @@ tf2onnx.tfonnx.process_tf_graph(tf_graph,
             inputs_as_nchw: transpose inputs in list from nchw to nchw
             input_names: list of input node names in graph, input name format as node_name:port_id
             output_names: list of output node names in graph, output name format as node_name:port_id
+            const_node_values: an optional dict mapping node names to tensor values
         Return:
             onnx graph
     """

tests/backend_test_base.py

@@ -26,6 +26,8 @@
 from tf2onnx import optimizer
 from tf2onnx.tf_loader import tf_reset_default_graph, tf_session, tf_placeholder, from_function, freeze_session
 from tf2onnx.tf_loader import tf_optimize, is_tf2
+from tf2onnx.tf_utils import compress_graph_def
+from tf2onnx.graph import ExternalTensorStorage
 
 
 class Tf2OnnxBackendTestBase(unittest.TestCase):
@@ -72,9 +74,10 @@ def run_onnxruntime(self, model_path, inputs, output_names):
         results = m.run(output_names, inputs)
         return results
 
-    def run_backend(self, g, outputs, input_dict):
-        model_proto = g.make_model("test")
-        model_path = self.save_onnx_model(model_proto, input_dict)
+    def run_backend(self, g, outputs, input_dict, large_model=False):
+        tensor_storage = ExternalTensorStorage() if large_model else None
+        model_proto = g.make_model("test", external_tensor_storage=tensor_storage)
+        model_path = self.save_onnx_model(model_proto, input_dict, external_tensor_storage=tensor_storage)
 
         if self.config.backend == "onnxruntime":
             y = self.run_onnxruntime(model_path, input_dict, outputs)
@@ -86,7 +89,8 @@ def run_backend(self, g, outputs, input_dict):
 
     def run_test_case(self, func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-07, atol=1e-5,
                       convert_var_to_const=True, constant_fold=True, check_value=True, check_shape=True,
-                      check_dtype=True, process_args=None, onnx_feed_dict=None, graph_validator=None, as_session=False):
+                      check_dtype=True, process_args=None, onnx_feed_dict=None, graph_validator=None, as_session=False,
+                      large_model=False):
         # optional - passed to process_tf_graph
         if process_args is None:
             process_args = {}
@@ -121,7 +125,9 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
             concrete_func = tf.function(func, input_signature=tuple(input_tensors))
             concrete_func = concrete_func.get_concrete_function()
             graph_def = from_function(concrete_func,
-                                      input_names=list(feed_dict.keys()), output_names=output_names_with_port)
+                                      input_names=list(feed_dict.keys()),
+                                      output_names=output_names_with_port,
+                                      large_model=large_model)
         else:
             #
             # use graph to execute the tensorflow func
@@ -151,6 +157,9 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
 
         tf_reset_default_graph()
         with tf_session() as sess:
+            const_node_values = None
+            if large_model:
+                const_node_values = compress_graph_def(graph_def)
             tf.import_graph_def(graph_def, name='')
 
             if self.config.is_debug_mode:
@@ -161,9 +170,11 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
             g = process_tf_graph(sess.graph, opset=self.config.opset,
                                  input_names=list(feed_dict.keys()),
                                  output_names=output_names_with_port,
-                                 target=self.config.target, **process_args)
+                                 target=self.config.target,
+                                 const_node_values=const_node_values,
+                                 **process_args)
             g = optimizer.optimize_graph(g)
-            actual = self.run_backend(g, output_names_with_port, onnx_feed_dict)
+            actual = self.run_backend(g, output_names_with_port, onnx_feed_dict, large_model)
 
         for expected_val, actual_val in zip(expected, actual):
             if check_value:
@@ -180,10 +191,11 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
 
         return g
 
-    def save_onnx_model(self, model_proto, feed_dict, postfix=""):
+    def save_onnx_model(self, model_proto, feed_dict, postfix="", external_tensor_storage=None):
         target_path = utils.save_onnx_model(self.test_data_directory, self._testMethodName + postfix, feed_dict,
                                             model_proto, include_test_data=self.config.is_debug_mode,
-                                            as_text=self.config.is_debug_mode)
+                                            as_text=self.config.is_debug_mode,
+                                            external_tensor_storage=external_tensor_storage)
 
         self.logger.debug("create model file: %s", target_path)
         return target_path

tests/test_backend.py

@@ -820,6 +820,15 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_tf_min_version("2.2")
+    def test_large_model_format(self):
+        x_val = np.array([2.0], dtype=np.float32)
+        y_const = np.arange(2000, dtype=np.float32)
+        def func(x):
+            x_ = tf.multiply(x, tf.constant(y_const))
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, large_model=True)
+
     @check_target('rs6', 'GatherNd')
     def test_gathernd(self):
         x_val = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)

tests/test_convert.py

@@ -8,7 +8,7 @@
 import unittest
 
 from tf2onnx import convert
-
+from common import check_tf_min_version
 
 def run_test_case(args):
     """ run case and clean up """
@@ -33,6 +33,18 @@ def test_convert_saved_model(self):
                                        '--output',
                                        'converted_saved_model.onnx']))
 
+    @check_tf_min_version("2.2")
+    def test_convert_large_model(self):
+        """ convert saved model to onnx large model format """
+        self.assertTrue(run_test_case(['',
+                                       '--large_model',
+                                       '--saved-model',
+                                       'tests/models/regression/saved_model',
+                                       '--tag',
+                                       'serve',
+                                       '--output',
+                                       'converted_saved_model.zip']))
+
     def test_convert_graphdef(self):
         """ convert graphdef """
         self.assertTrue(run_test_case(['',

tests/test_internals.py

@@ -139,7 +139,7 @@ def test_rewrite_subgraph(self):
             op_name = utils.make_name("ReplacedOp")
             out_name = utils.port_name(op_name)
             new_node = g.make_node("Sub", inputs=input_node.input, outputs=[out_name], name=op_name)
-            g.replace_all_inputs(ops, output_node.output[0], new_node.output[0])
+            g.replace_all_inputs(output_node.output[0], new_node.output[0])  # ops=ops
             for n in set(match.get_nodes()):
                 g.remove_node(n.name)
         g.topological_sort(ops)

tf2onnx/convert.py

@@ -22,6 +22,8 @@
 from tf2onnx.tfonnx import process_tf_graph
 from tf2onnx import constants, logging, utils, optimizer
 from tf2onnx import tf_loader
+from tf2onnx.graph import ExternalTensorStorage
+from tf2onnx.tf_utils import compress_graph_def
 
 # pylint: disable=unused-argument
 
@@ -53,6 +55,7 @@ def get_args():
                         help="For TF2.x saved_model, index of func signature in __call__ (--signature_def is ignored)")
     parser.add_argument("--checkpoint", help="input from checkpoint")
     parser.add_argument("--keras", help="input from keras model")
+    parser.add_argument("--large_model", help="use the large model format (for models > 2GB)", action="store_true")
     parser.add_argument("--output", help="output model file")
     parser.add_argument("--inputs", help="model input_names")
     parser.add_argument("--outputs", help="model output_names")
@@ -129,7 +132,8 @@ def main():
         model_path = args.checkpoint
     if args.saved_model:
         graph_def, inputs, outputs = tf_loader.from_saved_model(
-            args.saved_model, args.inputs, args.outputs, args.tag, args.signature_def, args.concrete_function)
+            args.saved_model, args.inputs, args.outputs, args.tag,
+            args.signature_def, args.concrete_function, args.large_model)
         model_path = args.saved_model
     if args.keras:
         graph_def, inputs, outputs = tf_loader.from_keras(
@@ -141,6 +145,9 @@ def main():
         logger.info("outputs: %s", outputs)
 
     with tf.Graph().as_default() as tf_graph:
+        const_node_values = None
+        if args.large_model:
+            const_node_values = compress_graph_def(graph_def)
         tf.import_graph_def(graph_def, name='')
     with tf_loader.tf_session(graph=tf_graph):
         g = process_tf_graph(tf_graph,
@@ -152,17 +159,24 @@ def main():
                              shape_override=args.shape_override,
                              input_names=inputs,
                              output_names=outputs,
-                             inputs_as_nchw=args.inputs_as_nchw)
+                             inputs_as_nchw=args.inputs_as_nchw,
+                             const_node_values=const_node_values)
 
     onnx_graph = optimizer.optimize_graph(g)
-    model_proto = onnx_graph.make_model("converted from {}".format(model_path))
+
+    tensor_storage = ExternalTensorStorage() if args.large_model else None
+    model_proto = onnx_graph.make_model("converted from {}".format(model_path), external_tensor_storage=tensor_storage)
 
     # write onnx graph
     logger.info("")
     logger.info("Successfully converted TensorFlow model %s to ONNX", model_path)
     if args.output:
-        utils.save_protobuf(args.output, model_proto)
-        logger.info("ONNX model is saved at %s", args.output)
+        if args.large_model:
+            utils.save_onnx_zip(args.output, model_proto, tensor_storage)
+            logger.info("Zipped ONNX model is saved at %s. Unzip before opening in onnxruntime.", args.output)
+        else:
+            utils.save_protobuf(args.output, model_proto)
+            logger.info("ONNX model is saved at %s", args.output)
     else:
         logger.info("To export ONNX model to file, please run with `--output` option")