(WIP) respond to PR comments

Author: phi-dbq

python/sparkdl/graph/input.py

@@ -29,43 +29,20 @@ class TFInputGraph(object):
     """
 
     # TODO: for (de-)serialization, the class should correspond to a ProtocolBuffer definition.
-    def __init__(self, graph_def, input_mapping, output_mapping):
+    def __init__(self, graph_def):
         # tf.GraphDef
         self.graph_def = graph_def
 
-        if isinstance(input_mapping, dict):
-            input_mapping = list(input_mapping.items())
-        assert isinstance(input_mapping, list), \
-            "output mapping must be a list of strings, found type {}".format(type(input_mapping))
-        self.input_mapping = sorted(input_mapping)
-
-        if isinstance(output_mapping, dict):
-            output_mapping = list(output_mapping.items())
-        assert isinstance(output_mapping, list), \
-            "output mapping must be a list of strings, found type {}".format(type(output_mapping))
-        self.output_mapping = sorted(output_mapping)
-
-
-def _get_params_from(gin_builder, input_mapping, output_mapping):
-    gin = gin_builder.build(input_mapping, output_mapping)
-    imap = dict(gin.input_mapping)
-    assert len(imap) == len(gin.input_mapping)
-    omap = dict(gin.output_mapping)
-    assert len(omap) == len(gin.output_mapping)
-    return gin.graph_def, imap, omap
-
-
 def get_params_from_checkpoint(checkpoint_dir, signature_def_key, input_mapping, output_mapping):
     assert signature_def_key is not None
     gin_builder = TFInputGraphBuilder.fromCheckpoint(checkpoint_dir, signature_def_key)
-    return _get_params_from(gin_builder, input_mapping, output_mapping)
-
+    return gin_builder.build(input_mapping, output_mapping)
 
 def get_params_from_saved_model(saved_model_dir, tag_set, signature_def_key, input_mapping,
                                 output_mapping):
     assert signature_def_key is not None
     gin_builder = TFInputGraphBuilder.fromSavedModel(saved_model_dir, tag_set, signature_def_key)
-    return _get_params_from(gin_builder, input_mapping, output_mapping)
+    return gin_builder.build(input_mapping, output_mapping)
 
 
 class TFInputGraphBuilder(object):
@@ -117,12 +94,19 @@ def build(self, input_mapping, output_mapping):
                     tnsr = tnsr_or_sig
                 fetches.append(tfx.get_tensor(graph, tnsr))
                 tf_output_colname = tfx.op_name(graph, tnsr)
+                # NOTE(phi-dbq): put the check here as it will be the entry point to construct
+                #                a `TFInputGraph` object.
+                assert tf_output_colname not in _output_mapping, \
+                    "operation {} has multiple output tensors and ".format(tf_output_colname) + \
+                    "at least two of them are used in the output DataFrame. " + \
+                    "Operation names are used to name columns which leads to conflicts. "  + \
+                    "You can apply `tf.identity` ops to each to avoid name conflicts."
                 _output_mapping[tf_output_colname] = requested_colname
             output_mapping = _output_mapping
 
             gdef = tfx.strip_and_freeze_until(fetches, graph, sess)
 
-        return TFInputGraph(gdef, input_mapping, output_mapping)
+        return TFInputGraph(gdef), input_mapping, output_mapping
 
     @classmethod
     def fromGraph(cls, graph):
@@ -167,6 +151,8 @@ def import_graph_fn(sess):
                 ckpt_path = tf.train.latest_checkpoint(checkpoint_dir)
 
                 # NOTE(phi-dbq): we must manually load meta_graph_def to get the signature_def
+                #                the current `import_graph_def` function seems to ignore
+                #                any signature_def fields in a checkpoint's meta_graph_def.
                 meta_graph_def = meta_graph_pb2.MetaGraphDef()
                 with open("{}.meta".format(ckpt_path), 'rb') as fin:
                     meta_graph_def.ParseFromString(fin.read())
@@ -177,7 +163,9 @@ def import_graph_fn(sess):
                 sig_def = None
                 if signature_def_key is not None:
                     sig_def = meta_graph_def.signature_def[signature_def_key]
-                    # TODO: check if sig_def is valid
+                    assert sig_def, 'singnature_def_key {} provided, '.format(signature_def_key) + \
+                        'but failed to find it from the meta_graph_def ' + \
+                        'from checkpoint {}'.format(checkpoint_dir)
 
             return sig_def
 

python/sparkdl/transformers/param.py

@@ -29,10 +29,10 @@
 from sparkdl.graph.builder import GraphFunction, IsolatedSession
 import sparkdl.graph.utils as tfx
 from sparkdl.graph.input import TFInputGraph, TFInputGraphBuilder
-"""
-Copied from PySpark for backward compatibility. First in Apache Spark version 2.1.1.
-"""
 
+########################################################
+# Copied from PySpark for backward compatibility. First in Apache Spark version 2.1.1.
+########################################################
 
 def keyword_only(func):
     """
@@ -98,34 +98,22 @@ def getOutputCol(self):
         return self.getOrDefault(self.outputCol)
 
 
-"""
-TensorFlow Specific Parameters
-"""
-
+########################################################
+# New in sparkdl: TensorFlow Specific Parameters
+########################################################
 
 class SparkDLTypeConverters(object):
     @staticmethod
     def toTFGraph(value):
         if isinstance(value, tf.Graph):
             return value
-        elif isinstance(value, GraphFunction):
-            with IsolatedSession() as issn:
-                issn.importGraphFunction(value, prefix='')
-                g = issn.graph
-            return g
         else:
             raise TypeError("Could not convert %s to TensorFlow Graph" % type(value))
 
     @staticmethod
     def toTFInputGraph(value):
         if isinstance(value, TFInputGraph):
             return value
-        elif isinstance(value, TFInputGraphBuilder):
-            return value
-        elif isinstance(value, tf.Graph):
-            return TFInputGraphBuilder.fromGraph(value)
-        elif isinstance(value, tf.GraphDef):
-            return TFInputGraphBuilder.fromGraphDef(value)
         else:
             raise TypeError("Could not convert %s to TFInputGraph" % type(value))
 
@@ -171,9 +159,6 @@ def converter(value):
         return converter
 
 
-# New in sparkdl
-
-
 class HasOutputMapping(Params):
     """
     Mixin for param outputMapping: ordered list of ('outputTensorOpName', 'outputColName') pairs
@@ -185,7 +170,11 @@ class HasOutputMapping(Params):
         typeConverter=SparkDLTypeConverters.asTensorToColumnMap)
 
     def setOutputMapping(self, value):
-        return self._set(outputMapping=value)
+        # NOTE(phi-dbq): due to the nature of TensorFlow import modes, we can only derive the
+        #                serializable TFInputGraph object once the inputMapping and outputMapping
+        #                parameters are provided.
+        raise NotImplementedError(
+            "Please use the Transformer's constructor to assign `outputMapping` field.")
 
     def getOutputMapping(self):
         return self.getOrDefault(self.outputMapping)
@@ -202,28 +191,36 @@ class HasInputMapping(Params):
         typeConverter=SparkDLTypeConverters.asColumnToTensorMap)
 
     def setInputMapping(self, value):
-        return self._set(inputMapping=value)
+        # NOTE(phi-dbq): due to the nature of TensorFlow import modes, we can only derive the
+        #                serializable TFInputGraph object once the inputMapping and outputMapping
+        #                parameters are provided.
+        raise NotImplementedError(
+            "Please use the Transformer's constructor to assigne `inputMapping` field.")
 
     def getInputMapping(self):
         return self.getOrDefault(self.inputMapping)
 
 
 class HasTFInputGraph(Params):
     """
-    Mixin for param tfGraph: the :class:`tf.Graph` object that represents a TensorFlow computation.
+    Mixin for param tfInputGraph: a serializable object derived from a TensorFlow computation graph.
     """
     tfInputGraph = Param(
         Params._dummy(),
         "tfInputGraph",
-        "TensorFlow Graph object",
+        "A serializable object derived from a TensorFlow computation graph",
         typeConverter=SparkDLTypeConverters.toTFInputGraph)
 
     def __init__(self):
         super(HasTFInputGraph, self).__init__()
         self._setDefault(tfInputGraph=None)
 
     def setTFInputGraph(self, value):
-        return self._set(tfInputGraph=value)
+        # NOTE(phi-dbq): due to the nature of TensorFlow import modes, we can only derive the
+        #                serializable TFInputGraph object once the inputMapping and outputMapping
+        #                parameters are provided.
+        raise NotImplementedError(
+            "Please use the Transformer's constructor to assign `tfInputGraph` field.")
 
     def getTFInputGraph(self):
         return self.getOrDefault(self.tfInputGraph)
@@ -236,7 +233,7 @@ class HasTFHParams(Params):
     tfHParams = Param(
         Params._dummy(),
         "hparams",
-        "instance of :class:`tf.contrib.training.HParams`",
+        "instance of :class:`tf.contrib.training.HParams`, a key-value map-like object",
         typeConverter=SparkDLTypeConverters.toTFHParams)
 
     def setTFHParams(self, value):

python/sparkdl/transformers/tf_tensor.py

@@ -21,7 +21,7 @@
 from pyspark.ml import Transformer
 
 import sparkdl.graph.utils as tfx
-from sparkdl.graph.input import TFInputGraphBuilder
+from sparkdl.graph.input import TFInputGraph, TFInputGraphBuilder
 from sparkdl.transformers.param import (keyword_only, SparkDLTypeConverters, HasInputMapping,
                                         HasOutputMapping, HasTFInputGraph, HasTFHParams)
 
@@ -57,15 +57,38 @@ def setParams(self, tfInputGraph=None, inputMapping=None, outputMapping=None, tf
         """
         super(TFTransformer, self).__init__()
         kwargs = self._input_kwargs
-        _maybe_gin = SparkDLTypeConverters.toTFInputGraph(tfInputGraph)
+        # The set of parameters either come from some helper functions,
+        # in which case type(_maybe_gin) is already TFInputGraph.
+        _maybe_gin = tfInputGraph
+        if isinstance(_maybe_gin, TFInputGraph):
+            return self._set(**kwargs)
+
+        # Otherwise, `_maybe_gin` needs to be converted to TFInputGraph
+        # We put all the conversion logic here rather than in SparkDLTypeConverters
         if isinstance(_maybe_gin, TFInputGraphBuilder):
-            kwargs['tfInputGraph'] = _maybe_gin.build(inputMapping, outputMapping)
+            gin = _maybe_gin
+        elif isinstance(_maybe_gin, tf.Graph):
+            gin = TFInputGraphBuilder.fromGraph(_maybe_gin)
+        elif isinstance(_maybe_gin, tf.GraphDef):
+            gin = TFInputGraphBuilder.fromGraphDef(_maybe_gin)
+        else:
+            raise TypeError("TFTransformer expect tfInputGraph convertible to TFInputGraph, " + \
+                            "but the given type {} cannot be converted, ".format(type(tfInputGraph)) + \
+                            "please provide `tf.Graph`, `tf.GraphDef` or use one of the " + \
+                            "`get_params_from_*` helper functions to build parameters")
+
+        gin, input_mapping, output_mapping = gin.build(inputMapping, outputMapping)
+        kwargs['tfInputGraph'] = gin
+        kwargs['inputMapping'] = input_mapping
+        kwargs['outputMapping'] = output_mapping
+
+        # Further conanonicalization, e.g. converting dict to sorted str pairs happens here
         return self._set(**kwargs)
 
     def _transform(self, dataset):
         gin = self.getTFInputGraph()
-        input_mapping = gin.input_mapping
-        output_mapping = gin.output_mapping
+        input_mapping = self.getInputMapping()
+        output_mapping = self.getOutputMapping()
 
         graph = tf.Graph()
         with tf.Session(graph=graph):

python/tests/transformers/tf_tensor_test.py

@@ -54,7 +54,6 @@ def test_build_from_tf_graph(self):
 
         # Build the TensorFlow graph
         with tf.Session() as sess:
-            #x = tf.placeholder(tf.float64, shape=[None, vec_size])
             x = tfs.block(analyzed_df, 'vec')
             z = tf.reduce_mean(x, axis=1)
             graph = sess.graph
@@ -68,9 +67,9 @@ def test_build_from_tf_graph(self):
 
         # Apply the transform
         gin_from_graph = TFInputGraphBuilder.fromGraph(graph)
-        for gin in [gin_from_graph, graph]:
+        for gin_or_graph in [gin_from_graph, graph]:
             transfomer = TFTransformer(
-                tfInputGraph=TFInputGraphBuilder.fromGraph(graph),
+                tfInputGraph=gin_or_graph,
                 inputMapping={
                     'vec': x
                 },