Update run_pretrained_models.py to support large models

tests/run_pretrained_models.py

@@ -41,6 +41,7 @@
 from tf2onnx import tf_loader, logging, optimizer, utils, tf_utils
 from tf2onnx.tfonnx import process_tf_graph
 from tf2onnx.tf_loader import tf_session, tf_reset_default_graph
+from tf2onnx.graph import ExternalTensorStorage
 
 logger = logging.getLogger("run_pretrained")
 
@@ -102,16 +103,20 @@ class Test(object):
     cache_dir = None
     target = []
 
-    def __init__(self, url, local, make_input, input_names, output_names,
+    def __init__(self, url, local, input_func, input_names, output_names,
                  disabled=False, rtol=0.01, atol=1e-6,
                  check_only_shape=False, model_type="frozen", force_input_shape=False,
-                 skip_tensorflow=False, opset_constraints=None, tf_min_version=None, tag=None):
+                 skip_tensorflow=False, opset_constraints=None, tf_min_version=None, tag=None,
+                 skip_conversion=False, converted_model=None, signature_def=None, concrete_function=None,
+                 large_model=False, structured_outputs=None):
         self.url = url
-        self.make_input = make_input
+        self.input_func = input_func
         self.local = local
         self.input_names = input_names
         self.output_names = output_names
         self.disabled = disabled
+        self.large_model = large_model
+        self.structured_outputs = structured_outputs  # Needed to determine output order for tf_function
         self.rtol = rtol
         self.atol = atol
         self.check_only_shape = check_only_shape
@@ -122,8 +127,18 @@ def __init__(self, url, local, make_input, input_names, output_names,
         self.tag = tag
         self.force_input_shape = force_input_shape
         self.skip_tensorflow = skip_tensorflow
+        self.skip_conversion = skip_conversion
+        self.converted_model = converted_model
         self.opset_constraints = opset_constraints
         self.tf_min_version = tf_min_version
+        self.signatures = [signature_def] if signature_def else None
+        self.concrete_function = concrete_function
+
+    def make_input(self, v):
+        """Allows each input to specify its own function while defaulting to the input_get function"""
+        if isinstance(v, dict):
+            return _INPUT_FUNC_MAPPING[v["input_get"]](v["shape"])
+        return self.input_func(v)
 
     def download_model(self):
         """Download model from url."""
@@ -149,7 +164,7 @@ def download_model(self):
         if not os.path.exists(fpath):
             utils.get_url(url, fpath)
         model_path = os.path.join(dir_name, self.local)
-        if not os.path.exists(model_path):
+        if not os.path.exists(model_path) or self.local == ".":
             if ftype == 'tgz':
                 tar = tarfile.open(fpath)
                 tar.extractall(dir_name)
@@ -179,19 +194,23 @@ def run_tensorflow(self, sess, inputs):
         for k, v in inputs.items():
             k = sess.graph.get_tensor_by_name(k)
             feed_dict[k] = v
+        logger.info("Running TF")
         result = sess.run(self.output_names, feed_dict=feed_dict)
         if self.perf:
+            logger.info("Running TF perf")
             start = time.time()
             for _ in range(PERFITER):
                 _ = sess.run(self.output_names, feed_dict=feed_dict)
             self.tf_runtime = time.time() - start
         return result
 
-    def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, input_names=None):
+    def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, input_names=None,
+                const_node_values=None):
         """Convert graph to tensorflow."""
         return process_tf_graph(tf_graph, continue_on_error=False, opset=opset,
                                 extra_opset=extra_opset, target=Test.target, shape_override=shape_override,
-                                input_names=input_names, output_names=self.output_names)
+                                input_names=input_names, output_names=self.output_names,
+                                const_node_values=const_node_values)
 
     def run_caffe2(self, name, model_proto, inputs):
         """Run test again caffe2 backend."""
@@ -205,11 +224,12 @@ def run_caffe2(self, name, model_proto, inputs):
             self.onnx_runtime = time.time() - start
         return results
 
-    def run_onnxruntime(self, name, model_proto, inputs):
+    def run_onnxruntime(self, name, model_proto, inputs, external_tensor_storage=None):
         """Run test against onnxruntime backend."""
         import onnxruntime as rt
         model_path = utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=True,
-                                           as_text=utils.is_debug_mode())
+                                           as_text=utils.is_debug_mode(),
+                                           external_tensor_storage=external_tensor_storage)
         logger.info("Model saved to %s", model_path)
         m = rt.InferenceSession(model_path)
         results = m.run(self.output_names, inputs)
@@ -221,10 +241,14 @@ def run_onnxruntime(self, name, model_proto, inputs):
         return results
 
     @staticmethod
-    def create_onnx_file(name, model_proto, inputs, outdir):
+    def create_onnx_file(name, model_proto, inputs, outdir, external_tensor_storage=None):
         os.makedirs(outdir, exist_ok=True)
-        model_path = os.path.join(outdir, name + ".onnx")
-        utils.save_protobuf(model_path, model_proto)
+        if external_tensor_storage is None:
+            model_path = os.path.join(outdir, name + ".onnx")
+            utils.save_protobuf(model_path, model_proto)
+        else:
+            model_path = os.path.join(outdir, name + ".zip")
+            utils.save_onnx_zip(model_path, model_proto, external_tensor_storage)
         logger.info("Created %s", model_path)
 
     def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_opset=None,
@@ -236,7 +260,7 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
         if self.url:
             _, dir_name = self.download_model()
             logger.info("Downloaded to %s", dir_name)
-            model_path = os.path.join(dir_name, self.local)
+            model_path = os.path.join(dir_name, self.local) if self.local != "." else dir_name
         else:
             model_path = self.local
 
@@ -246,13 +270,15 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
         if self.model_type in ["checkpoint"]:
             graph_def, input_names, outputs = tf_loader.from_checkpoint(model_path, input_names, outputs)
         elif self.model_type in ["saved_model"]:
-            try:
-                res = tf_loader.from_saved_model(model_path, input_names, outputs, self.tag)
-            except OSError:
-                model_path = dir_name
-                logger.info("Load model(2) from %r", model_path)
-                res = tf_loader.from_saved_model(model_path, input_names, outputs, self.tag)
-            graph_def, input_names, outputs = res[:3]
+            loaded = tf_loader.from_saved_model(model_path, input_names, outputs, self.tag, self.signatures,
+                                                self.concrete_function, self.large_model,
+                                                return_concrete_func=self.large_model)
+            if self.large_model:
+                # Must maintain ref to imported since concrete_func uses weak refs
+                # pylint: disable=unused-variable
+                graph_def, input_names, outputs, concrete_func, imported = loaded
+            else:
+                graph_def, input_names, outputs = loaded
         elif self.model_type in ["keras"]:
             graph_def, input_names, outputs = tf_loader.from_keras(model_path, input_names, outputs)
         else:
@@ -261,9 +287,34 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
         if utils.is_debug_mode():
             utils.save_protobuf(os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"), graph_def)
 
+        if self.large_model:
+            inputs = {}
+            for k in input_names:
+                v = self.input_names[k]
+                inputs[k.split(":")[0]] = tf.constant(self.make_input(v))
+            tf_func = tf.function(concrete_func)
+            logger.info("Running TF")
+            tf_results_d = tf_func(**inputs)
+            if self.structured_outputs is None:
+                tf_results = list(tf_results_d.values())
+            else:
+                tf_results = [tf_results_d[output] for output in self.structured_outputs]
+            if self.perf:
+                logger.info("Running TF perf")
+                start = time.time()
+                for _ in range(PERFITER):
+                    _ = concrete_func(**inputs)
+                self.tf_runtime = time.time() - start
+            logger.info("TensorFlow OK")
+
         inputs = {}
         shape_override = {}
         tf_reset_default_graph()
+
+        from tf2onnx.tf_utils import compress_graph_def
+        const_node_values = None
+        if self.large_model:
+            const_node_values = compress_graph_def(graph_def)
         g = tf.import_graph_def(graph_def, name='')
         # with tf_session(config=tf.ConfigProto(allow_soft_placement=True), graph=g) as sess:
         with tf_session(graph=g) as sess:
@@ -288,30 +339,50 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
             # run the model with tensorflow
             if self.skip_tensorflow:
                 logger.info("TensorFlow SKIPPED")
-            else:
+            elif not self.large_model:
                 tf_results = self.run_tensorflow(sess, inputs)
                 logger.info("TensorFlow OK")
 
         model_proto = None
-        try:
-            # convert model to onnx
-            onnx_graph = self.to_onnx(sess.graph, opset=opset, extra_opset=extra_opset,
-                                      shape_override=shape_override, input_names=inputs.keys())
-            onnx_graph = optimizer.optimize_graph(onnx_graph)
-            model_proto = onnx_graph.make_model("converted from tf2onnx")
-            logger.info("To_ONNX, OK")
-            if onnx_file:
-                self.create_onnx_file(name, model_proto, inputs, onnx_file)
-        except Exception:
-            logger.error("To_ONNX FAIL", exc_info=1)
-            return False
+        if self.skip_conversion:
+            if self.large_model:
+                external_tensor_storage = ExternalTensorStorage()
+                model_proto = utils.model_proto_from_zip(self.converted_model, external_tensor_storage)
+            else:
+                external_tensor_storage = None
+                model_proto = utils.model_proto_from_file(self.converted_model)
+            logger.info("ONNX loaded from file")
+        else:
+            try:
+                # convert model to onnx
+                onnx_graph = self.to_onnx(sess.graph, opset=opset, extra_opset=extra_opset,
+                                          shape_override=shape_override, input_names=inputs.keys(),
+                                          const_node_values=const_node_values)
+                onnx_graph = optimizer.optimize_graph(onnx_graph)
+                print("ONNX", onnx_graph.dump_node_statistics())
+                external_tensor_storage = ExternalTensorStorage() if self.large_model else None
+                model_proto = onnx_graph.make_model("converted from tf2onnx",
+                                                    external_tensor_storage=external_tensor_storage)
+                logger.info("To_ONNX, OK")
+                if onnx_file:
+                    self.create_onnx_file(name, model_proto, inputs, onnx_file, external_tensor_storage)
+                if self.converted_model:
+                    if self.large_model:
+                        utils.save_onnx_zip(self.converted_model, model_proto, external_tensor_storage)
+                    else:
+                        utils.save_protobuf(self.converted_model, model_proto)
+                    logger.info("Created %s", self.converted_model)
+
+            except Exception:
+                logger.error("To_ONNX FAIL", exc_info=1)
+                return False
 
         try:
             onnx_results = None
             if backend == "caffe2":
                 onnx_results = self.run_caffe2(name, model_proto, inputs)
             elif backend == "onnxruntime":
-                onnx_results = self.run_onnxruntime(name, model_proto, inputs)
+                onnx_results = self.run_onnxruntime(name, model_proto, inputs, external_tensor_storage)
             else:
                 raise ValueError("unknown backend")
             logger.info("Run_ONNX OK")
@@ -390,6 +461,7 @@ def get_args():
     parser.add_argument("--list", help="list tests", action="store_true")
     parser.add_argument("--onnx-file", help="create onnx file in directory")
     parser.add_argument("--perf", help="capture performance numbers")
+    parser.add_argument("--perfiter", type=int, default=PERFITER, help="number of inferences for perf testing")
     parser.add_argument("--fold_const", help="enable tf constant_folding transformation before conversion",
                         action="store_true")
     parser.add_argument("--include-disabled", help="include disabled tests", action="store_true")
@@ -447,8 +519,9 @@ def load_tests_from_yaml(path):
                 opset_constraints.append(c)
 
         kwargs = {}
-        for kw in ["rtol", "atol", "disabled", "check_only_shape", "model_type",
-                   "skip_tensorflow", "force_input_shape", "tf_min_version", "tag"]:
+        for kw in ["rtol", "atol", "disabled", "check_only_shape", "model_type", "concrete_function",
+                   "skip_tensorflow", "force_input_shape", "tf_min_version", "tag", "skip_conversion",
+                   "converted_model", "signature_def", "large_model", "structured_outputs"]:
             if settings.get(kw) is not None:
                 kwargs[kw] = settings[kw]
 
@@ -459,6 +532,7 @@ def load_tests_from_yaml(path):
 
 
 def main():
+    global PERFITER
     args = get_args()
     logging.basicConfig(level=logging.get_verbosity_level(args.verbose))
     if args.debug:
@@ -477,6 +551,7 @@ def main():
 
     failed = 0
     count = 0
+    PERFITER = args.perfiter
     for test in test_keys:
         logger.info("===================================")
 
@@ -520,7 +595,8 @@ def main():
             for test in test_keys:
                 t = tests[test]
                 if t.perf:
-                    f.write("{},{},{}\n".format(test, t.tf_runtime, t.onnx_runtime))
+                    # Report perf in ms per inference
+                    f.write("{},{},{}\n".format(test, t.tf_runtime * 1000 / PERFITER, t.onnx_runtime * 1000 / PERFITER))
     return failed
 
 

tests/run_pretrained_models.yaml

@@ -119,7 +119,7 @@ benchtf-gru:
 esrgan-tf2:
   # url: https://tfhub.dev/captain-pool/esrgan-tf2/1/esrgan-tf2_1.tar.gz
   url: https://github.com/captain-pool/GSOC/releases/download/1.0.0/esrgan.tar.gz
-  model: ersgan
+  model: "."
   model_type: saved_model
   input_get: get_beach
   opset_constraints:

tf2onnx/tf_loader.py

@@ -316,18 +316,21 @@ def _from_saved_model_v2(model_path, input_names, output_names, tag, signature_d
             raise ValueError(err_large_model)
         raise e
 
-    return frozen_graph, inputs, outputs
+    return frozen_graph, inputs, outputs, concrete_func, imported
 
 
 def from_saved_model(model_path, input_names, output_names, tag=None,
-                     signatures=None, concrete_function=None, large_model=False):
+                     signatures=None, concrete_function=None, large_model=False, return_concrete_func=False):
     """Load tensorflow graph from saved_model."""
     if signatures is None:
         signatures = []
     tf_reset_default_graph()
     if is_tf2():
-        frozen_graph, input_names, output_names = \
+        frozen_graph, input_names, output_names, concrete_func, imported = \
             _from_saved_model_v2(model_path, input_names, output_names, tag, signatures, concrete_function, large_model)
+        if return_concrete_func:
+            tf_reset_default_graph()
+            return frozen_graph, input_names, output_names, concrete_func, imported
     else:
         with tf_session() as sess:
             frozen_graph, input_names, output_names = \

tf2onnx/utils.py

@@ -20,7 +20,7 @@
 from urllib3.util.retry import Retry
 import numpy as np
 from google.protobuf import text_format
-from onnx import helper, onnx_pb, defs, numpy_helper, __version__
+from onnx import helper, onnx_pb, defs, numpy_helper, ModelProto, __version__
 
 from . import constants
 
@@ -269,6 +269,22 @@ def save_protobuf(path, message, as_text=False):
         with open(path, "wb") as f:
             f.write(message.SerializeToString())
 
+def model_proto_from_file(model_path):
+    model_proto = ModelProto()
+    with open(model_path, "rb") as f:
+        model_proto.ParseFromString(f.read())
+    return model_proto
+
+def model_proto_from_zip(zip_path, external_tensor_storage):
+    model_proto = ModelProto()
+    with zipfile.ZipFile(zip_path, 'r') as z:
+        for n in z.namelist():
+            f = z.open(n)
+            if n.endswith(".onnx"):
+                model_proto.ParseFromString(f.read())
+            else:
+                external_tensor_storage.name_to_tensor_data[n] = f.read()
+    return model_proto
 
 def is_list_or_tuple(obj):
     return isinstance(obj, (list, tuple))