more tests

Author: phi-dbq

python/tests/transformers/tf_tensor_test.py

@@ -13,16 +13,24 @@
 # limitations under the License.
 #
 
+from keras.layers import Conv1D, Dense, Flatten, MaxPool1D
 import numpy as np
 import tensorflow as tf
+import tensorframes as tfs
 
 from pyspark.sql.types import Row
 
 from sparkdl.graph.builder import IsolatedSession
+import sparkdl.graph.utils as tfx
 from sparkdl.transformers.tf_tensor import TFOneDimTensorTransformer
 
 from ..tests import SparkDLTestCase
 
+def grab_df_arr(df, output_col):
+    """ Stack the numpy array from a DataFrame column """
+    return np.array([row.asDict()[output_col] 
+                     for row in df.select(output_col).toLocalIterator()])
+
 class TFOneDimTransformerTest(SparkDLTestCase):
 
     def test_simple(self):
@@ -52,3 +60,56 @@ def test_simple(self):
 
         out_tgt = np.array([row.outCol for row in final_df.select('outCol').collect()])
         self.assertTrue(np.allclose(out_ref, out_tgt))
+
+    
+    def test_map_blocks_graph(self):
+
+        vec_size = 17
+        num_vecs = 137
+
+        input_col = 'vec'
+        output_col = 'outCol'
+
+        df = self.session.createDataFrame([
+            Row(idx=idx, vec=np.random.randn(vec_size).tolist())
+            for idx in range(num_vecs)])
+        analyzed_df = tfs.analyze(df)
+
+        # Build the graph: the output should have the same leading/batch dimension
+        with IsolatedSession(using_keras=True) as issn:
+            tnsr_in = tfs.block(analyzed_df, input_col)
+            inp = tf.expand_dims(tnsr_in, axis=2)
+            inp = tf.cast(inp, tf.float32)
+            conv = Conv1D(filters=4, kernel_size=2)(inp)
+            pool = MaxPool1D(pool_size=2)(conv)
+            flat = Flatten()(pool)
+            dense = Dense(1)(flat)
+            redsum = tf.reduce_sum(dense, axis=1)
+            tnsr_out = tf.cast(redsum, tf.double, name='TnsrOut')
+
+            # Initialize the variables
+            init_op = tf.global_variables_initializer()
+            issn.run(init_op)
+            # Train the model ...
+            gfn = issn.asGraphFunction([tnsr_in], [tnsr_out])
+
+        
+        with IsolatedSession() as issn:
+            feeds, fetches = issn.importGraphFunction(gfn, prefix='')
+            orig_in_name = tfx.op_name(issn.graph, feeds[0])
+            input_df = analyzed_df.withColumnRenamed(input_col, orig_in_name)
+            output_df = tfs.map_blocks(fetches, input_df)
+            orig_out_name = tfx.op_name(issn.graph, fetches[0])
+            final_df = output_df.withColumnRenamed(orig_out_name, output_col)
+            
+        arr_ref = grab_df_arr(final_df, output_col)
+
+        # Using the Transformer
+        transformer = TFOneDimTensorTransformer(
+            graphFunction=gfn, inputCol=input_col, outputCol=output_col)
+        transformed_df = transformer.transform(analyzed_df)
+        
+        arr_tgt = grab_df_arr(transformed_df, output_col)
+
+        self.assertTrue(np.allclose(arr_ref, arr_tgt))
+