TensorFlow Transformer Part-2

python/sparkdl/graph/utils.py

@@ -16,8 +16,6 @@
 
 import logging
 import six
-import webbrowser
-from tempfile import NamedTemporaryFile
 
 import tensorflow as tf
 
@@ -95,31 +93,49 @@ def get_tensor(graph, tfobj_or_name):
         'cannot locate tensor {} in current graph'.format(_tensor_name)
     return tnsr
 
-def as_tensor_name(name):
+def as_tensor_name(tfobj_or_name):
     """
     Derive tf.Tensor name from an op/tensor name.
-    We do not check if the tensor exist (as no graph parameter is passed in).
+    If the input is a name, we do not check if the tensor exist
+    (as no graph parameter is passed in).
 
-    :param name: op name or tensor name
+    :param tfobj_or_name: either a tf.Tensor, tf.Operation or a name to either
     """
-    assert isinstance(name, six.string_types)
-    name_parts = name.split(":")
-    assert len(name_parts) <= 2, name_parts
-    if len(name_parts) < 2:
-        name += ":0"
-    return name
+    if isinstance(tfobj_or_name, six.string_types):
+        # If input is a string, assume it is a name and infer the corresponding tensor name.
+        # WARNING: this depends on TensorFlow's tensor naming convention
+        name = tfobj_or_name
+        name_parts = name.split(":")
+        assert len(name_parts) <= 2, name_parts
+        if len(name_parts) < 2:
+            name += ":0"
+        return name
+    elif hasattr(tfobj_or_name, 'graph'):
+        tfobj = tfobj_or_name
+        return get_tensor(tfobj.graph, tfobj).name
+    else:
+        raise TypeError('invalid tf.Tensor name query type {}'.format(type(tfobj_or_name)))
 
-def as_op_name(name):
+def as_op_name(tfobj_or_name):
     """
-    Derive tf.Operation name from an op/tensor name
-    We do not check if the operation exist (as no graph parameter is passed in).
+    Derive tf.Operation name from an op/tensor name.
+    If the input is a name, we do not check if the operation exist
+    (as no graph parameter is passed in).
 
-    :param name: op name or tensor name
+    :param tfobj_or_name: either a tf.Tensor, tf.Operation or a name to either
     """
-    assert isinstance(name, six.string_types)
-    name_parts = name.split(":")
-    assert len(name_parts) <= 2, name_parts
-    return name_parts[0]
+    if isinstance(tfobj_or_name, six.string_types):
+        # If input is a string, assume it is a name and infer the corresponding operation name.
+        # WARNING: this depends on TensorFlow's operation naming convention
+        name = tfobj_or_name
+        name_parts = name.split(":")
+        assert len(name_parts) <= 2, name_parts
+        return name_parts[0]
+    elif hasattr(tfobj_or_name, 'graph'):
+        tfobj = tfobj_or_name
+        return get_op(tfobj.graph, tfobj).name
+    else:
+        raise TypeError('invalid tf.Operation name query type {}'.format(type(tfobj_or_name)))
 
 def op_name(graph, tfobj_or_name):
     """

python/tests/graph/test_utils.py

@@ -0,0 +1,49 @@
+# Copyright 2017 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import tensorflow as tf
+
+import sparkdl.graph.utils as tfx
+
+from ..tests import PythonUnitTestCase
+
+class TFeXtensionGraphUtilsTest(PythonUnitTestCase):
+
+    def test_infer_graph_element_names(self):
+        for tnsr_idx in range(17):
+            op_name = 'someOp'
+            tnsr_name = '{}:{}'.format(op_name, tnsr_idx)
+            self.assertEqual(op_name, tfx.as_op_name(tnsr_name))
+            self.assertEqual(tnsr_name, tfx.as_tensor_name(tnsr_name))
+
+        with self.assertRaises(TypeError):
+            for wrong_value in [7, 1.2, tf.Graph()]:
+                tfx.as_op_name(wrong_value)
+                tfx.as_tensor_name(wrong_value)
+
+    def test_get_graph_elements(self):
+        op_name = 'someConstOp'
+        tnsr_name = '{}:0'.format(op_name)
+        tnsr = tf.constant(1427.08, name=op_name)
+        graph = tnsr.graph
+
+        self.assertEqual(op_name, tfx.as_op_name(tnsr))
+        self.assertEqual(op_name, tfx.op_name(graph, tnsr))
+        self.assertEqual(tnsr_name, tfx.as_tensor_name(tnsr))
+        self.assertEqual(tnsr_name, tfx.tensor_name(graph, tnsr))
+        self.assertEqual(tnsr, tfx.get_tensor(graph, tnsr))
+        self.assertEqual(tnsr.op, tfx.get_op(graph, tnsr))
+        self.assertEqual(graph, tfx.get_op(graph, tnsr).graph)
+        self.assertEqual(graph, tfx.get_tensor(graph, tnsr).graph)

python/tests/tests.py

@@ -29,6 +29,9 @@
 from pyspark.sql import SQLContext
 from pyspark.sql import SparkSession
 
+class PythonUnitTestCase(unittest.TestCase):
+    # Just the plain test unittest.TestCase, but won't have to do import check
+    pass
 
 class SparkDLTestCase(unittest.TestCase):