Merge pull request #984 from jignparm/jignparm/savedmodelv2

Add support for TF2.x saved_models from TFHub, as well as --tag & -concrete_function  cmd line parameters

Author: jignparm

README.md

@@ -139,6 +139,9 @@ python -m tf2onnx.convert
     [--outputs GRAPH_OUTPUS]
     [--inputs-as-nchw inputs_provided_as_nchw]
     [--opset OPSET]
+    [--tag TAG]
+    [--signature_def SIGNATURE_DEF]
+    [--concrete_function CONCRETE_FUNCTION]
     [--target TARGET]
     [--custom-ops list-of-custom-ops]
     [--fold_const]
@@ -176,6 +179,20 @@ By default we preserve the image format of inputs (`nchw` or `nhwc`) as given in
 
 By default we use the opset 8 to generate the graph. By specifying ```--opset``` the user can override the default to generate a graph with the desired opset. For example ```--opset 5``` would create a onnx graph that uses only ops available in opset 5. Because older opsets have in most cases fewer ops, some models might not convert on a older opset.
 
+#### --tag
+
+Only valid with parameter `--saved_model`. Specifies the tag in the saved_model to be used. Typical value is 'serve'.
+
+#### --signature_def
+
+Only valid with parameter `--saved_model`. Specifies which signature to use within the specified --tag value. Typical value is 'serving_default'.
+
+#### --concrete_function
+
+(This is experimental, valid only for TF2.x models)
+
+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.
+
 #### --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.

tests/run_pretrained_models.py

@@ -79,13 +79,18 @@ def get_ones(shape):
     """Get ones."""
     return np.ones(shape).astype(np.float32)
 
+def get_zeros(shape):
+    """Get zeros."""
+    return np.zeros(shape).astype(np.float32)
+
 
 _INPUT_FUNC_MAPPING = {
     "get_beach": get_beach,
     "get_random": get_random,
     "get_random256": get_random256,
     "get_ramp": get_ramp,
-    "get_ones": get_ones
+    "get_ones": get_ones,
+    "get_zeros": get_zeros,
 }
 
 OpsetConstraint = namedtuple("OpsetConstraint", "domain, min_version, max_version, excluded_version")
@@ -100,7 +105,7 @@ class Test(object):
     def __init__(self, url, local, make_input, 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):
+                 skip_tensorflow=False, opset_constraints=None, tf_min_version=None, tag=None):
         self.url = url
         self.make_input = make_input
         self.local = local
@@ -114,6 +119,7 @@ def __init__(self, url, local, make_input, input_names, output_names,
         self.tf_runtime = 0
         self.onnx_runtime = 0
         self.model_type = model_type
+        self.tag = tag
         self.force_input_shape = force_input_shape
         self.skip_tensorflow = skip_tensorflow
         self.opset_constraints = opset_constraints
@@ -240,7 +246,7 @@ 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"]:
-            graph_def, input_names, outputs = tf_loader.from_saved_model(model_path, input_names, outputs)
+            graph_def, input_names, outputs = tf_loader.from_saved_model(model_path, input_names, outputs, self.tag)
         elif self.model_type in ["keras"]:
             graph_def, input_names, outputs = tf_loader.from_keras(model_path, input_names, outputs)
         else:
@@ -436,7 +442,7 @@ def load_tests_from_yaml(path):
 
         kwargs = {}
         for kw in ["rtol", "atol", "disabled", "check_only_shape", "model_type",
-                   "skip_tensorflow", "force_input_shape", "tf_min_version"]:
+                   "skip_tensorflow", "force_input_shape", "tf_min_version", "tag"]:
             if settings.get(kw) is not None:
                 kwargs[kw] = settings[kw]
 

tests/run_pretrained_models.yaml

@@ -21,6 +21,7 @@ regression-checkpoint:
 regression-saved-model:
   model: models/regression/saved_model
   model_type: saved_model
+  tag: serve
   input_get: get_ramp
   inputs:
     "X:0": [1]
@@ -239,9 +240,10 @@ vgg-16:
 
 resnet50_v2_nchw: # NOTE: Tensorflow 1.9.0 fails
   skip_tensorflow: true # tensorflow fails: Default MaxPoolingOp only supports NHWC on device type CPU
-  model_type: saved_model
   url: http://download.tensorflow.org/models/official/20181001_resnet/savedmodels/resnet_v2_fp32_savedmodel_NCHW.tar.gz
   model: resnet_v2_fp32_savedmodel_NCHW/1538687196
+  model_type: saved_model
+  tag: serve
   input_get: get_beach
   inputs:
     "input_tensor:0": [64, 224, 224, 3]
@@ -250,9 +252,10 @@ resnet50_v2_nchw: # NOTE: Tensorflow 1.9.0 fails
     - softmax_tensor:0
 
 resnet50_v2_nhwc:
-  model_type: saved_model
   url: http://download.tensorflow.org/models/official/20181001_resnet/savedmodels/resnet_v2_fp32_savedmodel_NHWC.tar.gz
   model: resnet_v2_fp32_savedmodel_NHWC/1538687283
+  model_type: saved_model
+  tag: serve
   input_get: get_beach
   inputs:
     "input_tensor:0": [64, 224, 224, 3]

tests/test_convert.py

@@ -28,6 +28,8 @@ def test_convert_saved_model(self):
         self.assertTrue(run_test_case(['',
                                        '--saved-model',
                                        'tests/models/regression/saved_model',
+                                       '--tag',
+                                       'serve',
                                        '--output',
                                        'converted_saved_model.onnx']))
 

tf2onnx/convert.py

@@ -23,7 +23,6 @@
 from tf2onnx import constants, logging, utils, optimizer
 from tf2onnx import tf_loader
 
-
 # pylint: disable=unused-argument
 
 _HELP_TEXT = """
@@ -48,7 +47,10 @@ def get_args():
     parser.add_argument("--input", help="input from graphdef")
     parser.add_argument("--graphdef", help="input from graphdef")
     parser.add_argument("--saved-model", help="input from saved model")
-    parser.add_argument("--signature_def", help="signature_def from saved model to use")
+    parser.add_argument("--tag", help="tag to use for saved_model")
+    parser.add_argument("--signature_def", help="signature_def from saved_model to use")
+    parser.add_argument("--concrete_function", type=int, default=None,
+                        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("--output", help="output model file")
@@ -127,7 +129,7 @@ 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.signature_def)
+            args.saved_model, args.inputs, args.outputs, args.tag, args.signature_def, args.concrete_function)
         model_path = args.saved_model
     if args.keras:
         graph_def, inputs, outputs = tf_loader.from_keras(

tf2onnx/tf_loader.py

@@ -178,10 +178,16 @@ def from_checkpoint(model_path, input_names, output_names):
     return frozen_graph, input_names, output_names
 
 
-def _from_saved_model_v1(sess, model_path, input_names, output_names, signatures):
+def _from_saved_model_v1(sess, model_path, input_names, output_names, tag, signatures):
     """Load tensorflow graph from saved_model."""
 
-    imported = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], model_path)
+    if tag is None:
+        tag = [tf.saved_model.tag_constants.SERVING]
+
+    if not isinstance(tag, list):
+        tag = [tag]
+
+    imported = tf.saved_model.loader.load(sess, tag, model_path)
     for k in imported.signature_def.keys():
         if k.startswith("_"):
             # consider signatures starting with '_' private
@@ -209,43 +215,67 @@ def _from_saved_model_v1(sess, model_path, input_names, output_names, signatures
     return frozen_graph, input_names, output_names
 
 
-def _from_saved_model_v2(model_path, input_names, output_names, signatures):
+def _from_saved_model_v2(model_path, input_names, output_names, tag, signature_def, concrete_function_index):
     """Load tensorflow graph from saved_model."""
-    imported = tf.saved_model.load(model_path)  # pylint: disable=no-value-for-parameter
 
-    # f = meta_graph_def.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
-    for k in imported.signatures.keys():
-        if k.startswith("_"):
-            # consider signatures starting with '_' private
-            continue
-        signatures.append(k)
-    for k in signatures:
-        concrete_func = imported.signatures[k]
-        input_names = [input_tensor.name for input_tensor in concrete_func.inputs
-                       if input_tensor.dtype != tf.dtypes.resource]
-        output_names = [output_tensor.name for output_tensor in concrete_func.outputs
-                        if output_tensor.dtype != tf.dtypes.resource]
+    wrn_no_tag = "'--tag' not specified for saved_model. Using empty tag [[]]"
+    wrn_sig_1 = "'--signature_def' not specified, using first signature: %s"
+    err_many_sig = "Cannot load multiple signature defs in TF2.x: %s"
+    err_no_call = "Model doesn't contain usable concrete functions under  __call__. Try --signature-def instead."
+    err_index = "Invalid concrete_function value: %i. Valid values are [0 to %i]"
+    err_no_sig = "No signatures found in model. Try --concrete_function instead."
+    err_sig_nomatch = "Specified signature not in model %s"
+
+    if tag is None:
+        tag = [[]]
+        logger.warning(wrn_no_tag)
+    utils.make_sure(len(signature_def) < 2, err_many_sig, str(signature_def))
+    imported = tf.saved_model.load(model_path, tags=tag)  # pylint: disable=no-value-for-parameter
+
+    all_sigs = imported.signatures.keys()
+    valid_sigs = [s for s in all_sigs if not s.startswith("_")]
+    logger.info("Signatures found in model: %s", "[" + ",".join(valid_sigs) + "].")
+
+    concrete_func = None
+    if concrete_function_index is not None:
+        utils.make_sure(hasattr(imported, "__call__"), err_no_call)
+        utils.make_sure(concrete_function_index < len(imported.__call__.concrete_functions),
+                        err_index, concrete_function_index, len(imported.__call__.concrete_functions) - 1)
+        sig = imported.__call__.concrete_functions[concrete_function_index].structured_input_signature[0][0]
+        concrete_func = imported.__call__.get_concrete_function(sig)
+    elif signature_def:
+        utils.make_sure(signature_def[0] in valid_sigs, err_sig_nomatch, signature_def[0])
+        concrete_func = imported.signatures[signature_def[0]]
+    else:
+        utils.make_sure(len(valid_sigs) > 0, err_no_sig)
+        logger.warning(wrn_sig_1, valid_sigs[0])
+        concrete_func = imported.signatures[valid_sigs[0]]
 
-    frozen_graph = from_function(concrete_func, input_names, output_names)
-    return frozen_graph, input_names, output_names
+    inputs = [tensor.name for tensor in concrete_func.inputs if tensor.dtype != tf.dtypes.resource]
+    outputs = [tensor.name for tensor in concrete_func.outputs if tensor.dtype != tf.dtypes.resource]
 
+    # filter by user specified inputs/outputs
+    if input_names:
+        inputs = list(set(input_names) & set(inputs))
+    if output_names:
+        outputs = list(set(output_names) & set(outputs))
 
-def from_saved_model(model_path, input_names, output_names, signatures=None):
+    frozen_graph = from_function(concrete_func, inputs, outputs)
+    return frozen_graph, inputs, outputs
+
+
+def from_saved_model(model_path, input_names, output_names, tag=None, signatures=None, concrete_function=None):
     """Load tensorflow graph from saved_model."""
     if signatures is None:
         signatures = []
     tf_reset_default_graph()
     if is_tf2():
         frozen_graph, input_names, output_names = \
-            _from_saved_model_v2(model_path, input_names, output_names, signatures)
+            _from_saved_model_v2(model_path, input_names, output_names, tag, signatures, concrete_function)
     else:
         with tf_session() as sess:
             frozen_graph, input_names, output_names = \
-                _from_saved_model_v1(sess, model_path, input_names, output_names, signatures)
-
-    if len(signatures) > 1:
-        logger.warning("found multiple signatures %s in saved_model, pass --signature_def in command line",
-                       signatures)
+                _from_saved_model_v1(sess, model_path, input_names, output_names, tag, signatures)
 
     tf_reset_default_graph()
     return frozen_graph, input_names, output_names
@@ -366,6 +396,7 @@ def is_function(g):
         return 'tensorflow.python.framework.func_graph.FuncGraph' in str(type(g))
     return False
 
+
 _FUNCTIONS = {}