Merge pull request #53 from onnx/gs/onnx-1.2

add support for placeholderwithdefault, allow testing of checkpoints

Author: guschmue

tests/run_pretrained_models.py

@@ -17,6 +17,7 @@
 import yaml
 from tensorflow.core.framework import graph_pb2
 from tf2onnx.tfonnx import process_tf_graph
+from tensorflow.python.framework.graph_util import convert_variables_to_constants
 
 TMPPATH = tempfile.mkdtemp()
 PERFITER = 1000
@@ -77,11 +78,29 @@ def node_name(name):
     return name
 
 
+def freeze_session(sess, keep_var_names=None, output_names=None, clear_devices=True):
+    """Freezes the state of a session into a pruned computation graph."""
+    output_names = [i.replace(":0", "") for i in output_names]
+    graph = sess.graph
+    with graph.as_default():
+        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
+        output_names = output_names or []
+        output_names += [v.op.name for v in tf.global_variables()]
+        input_graph_def = graph.as_graph_def()
+        if clear_devices:
+            for node in input_graph_def.node:
+                node.device = ""
+        frozen_graph = convert_variables_to_constants(sess, input_graph_def,
+                                                      output_names, freeze_var_names)
+        return frozen_graph
+
+
 class Test(object):
     cache_dir = None
 
     def __init__(self, url, local, make_input, input_names, output_names,
-                 disabled=False, more_inputs=None, rtol=0.01, atol=0., check_only_shape=False):
+                 disabled=False, more_inputs=None, rtol=0.01, atol=0.,
+                 check_only_shape=False, model_type="frozen"):
         self.url = url
         self.make_input = make_input
         self.local = local
@@ -95,6 +114,7 @@ def __init__(self, url, local, make_input, input_names, output_names,
         self.perf = None
         self.tf_runtime = 0
         self.onnx_runtime = 0
+        self.model_type = model_type
 
     def download_file(self):
         """Download file from url."""
@@ -231,8 +251,18 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
             model_path = os.path.join(dir_name, self.local)
         else:
             model_path = self.local
+            dir_name = os.path.dirname(self.local)
         print("\tdownloaded", model_path)
 
+        # if the input model is a checkpoint, convert it to a frozen model
+        if self.model_type in ["checkpoint"]:
+            saver = tf.train.import_meta_graph(model_path)
+            with tf.Session() as sess:
+                saver.restore(sess, model_path[:-5])
+                frozen_graph = freeze_session(sess, output_names=self.output_names)
+                tf.train.write_graph(frozen_graph, dir_name, "frozen.pb", as_text=False)
+            model_path = os.path.join(dir_name, "frozen.pb")
+
         inputs = self.make_input(self.input_names)
         if self.more_inputs:
             for k, v in self.more_inputs.items():
@@ -314,7 +344,7 @@ def tests_from_yaml(fname):
         input_func = v.get("input_get")
         input_func = _INPUT_FUNC_MAPPING[input_func]
         kwargs = {}
-        for kw in ["rtol", "atol", "disabled", "more_inputs", "check_only_shape"]:
+        for kw in ["rtol", "atol", "disabled", "more_inputs", "check_only_shape", "model_type"]:
             if v.get(kw) is not None:
                 kwargs[kw] = v[kw]
 

tests/test_backend.py

@@ -362,6 +362,21 @@ def test_add(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
+    def test_placeholder(self):
+        x_val = np.array([1.0, 2.0, -3.0, -4.0], dtype=np.float32).reshape((2, 2))
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        output = tf.identity(x, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
+    def test_placeholder_with_default(self):
+        x_val = np.array([1.0, 2.0, -3.0, -4.0], dtype=np.float32).reshape((2, 2))
+        y = tf.constant(x_val, name="y")
+        x = tf.placeholder_with_default(y, x_val.shape, name=_TFINPUT)
+        output = tf.identity(x, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
     def test_add_bcast(self):
         x1_val = np.array([1.0, 2.0, -3.0, -4.0], dtype=np.float32).reshape((2, 2))
         x2_val = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], dtype=np.float32).reshape((2, 2, 2))

tf2onnx/graph.py

@@ -151,6 +151,8 @@ def get_tensor_value(self):
 
     def get_tensor(self):
         if not self.is_const():
+            if self.type == "Identity":
+                return self.inputs[0].get_tensor()
             raise ValueError("get tensor: {} must be Const".format(self.name))
         t = self.get_attr("value")
         if t:

tf2onnx/tfonnx.py

@@ -887,6 +887,7 @@ def minmax_op(ctx, node, name, args):
     "Pad": (pad_op, []),
     "Placeholder": (placeholder_op, []),
     "PlaceholderV2": (placeholder_op, []),
+    "PlaceholderWithDefault": (placeholder_op, []),
     "Pow": (pow_op, []),
     "Prod": (reduce_op, ["ReduceProd"]),
     "RandomNormal": (direct_op, []),