Fix context in setup: neuraxle===0.5.2 breaking changes

Author: alexbrillant

README.md

@@ -13,7 +13,7 @@ Neuraxle is a Machine Learning (ML) library for building neat pipelines, providi
 Create a tensorflow 1 model step by giving it a graph, an optimizer, and a loss function. 
 
 ```python
-def create_graph(step: TensorflowV1ModelStep):
+def create_graph(step: TensorflowV1ModelStep, context: ExecutionContext):
     tf.placeholder('float', name='data_inputs')
     tf.placeholder('float', name='expected_outputs')
 
@@ -24,7 +24,7 @@ def create_graph(step: TensorflowV1ModelStep):
 
 """
 # Note: you can also return a tuple containing two elements : tensor for training (fit), tensor for inference (transform)
-def create_graph(step: TensorflowV1ModelStep)
+def create_graph(step: TensorflowV1ModelStep, context: ExecutionContext)
     # ...
     decoder_outputs_training = create_training_decoder(step, encoder_state, decoder_cell)
     decoder_outputs_inference = create_inference_decoder(step, encoder_state, decoder_cell)
@@ -33,10 +33,10 @@ def create_graph(step: TensorflowV1ModelStep)
 """
 
 
-def create_loss(step: TensorflowV1ModelStep):
+def create_loss(step: TensorflowV1ModelStep, context: ExecutionContext):
     return tf.reduce_sum(tf.pow(step['output'] - step['expected_outputs'], 2)) / (2 * N_SAMPLES)
 
-def create_optimizer(step: TensorflowV1ModelStep):
+def create_optimizer(step: TensorflowV1ModelStep, context: ExecutionContext):
     return tf.train.GradientDescentOptimizer(step.hyperparams['learning_rate'])
 
 model_step = TensorflowV1ModelStep(
@@ -56,10 +56,10 @@ model_step = TensorflowV1ModelStep(
 Create a tensorflow 2 model step by giving it a model, an optimizer, and a loss function. 
 
 ```python
-def create_model(step: Tensorflow2ModelStep):
+def create_model(step: Tensorflow2ModelStep, context: ExecutionContext):
     return LinearModel()
 
-def create_optimizer(step: Tensorflow2ModelStep):
+def create_optimizer(step: Tensorflow2ModelStep, context: ExecutionContext):
     return tf.keras.optimizers.Adam(0.1)
 
 def create_loss(step: Tensorflow2ModelStep, expected_outputs, predicted_outputs):
@@ -94,36 +94,18 @@ feature_0_metric = metric_3d_to_2d_wrapper(mean_squared_error)
 metrics = {'mse': feature_0_metric}
 
 signal_prediction_pipeline = Pipeline([
-    ForEachDataInput(MeanStdNormalizer()),
-    ToNumpy(),
-    Tensorflow2ModelStep(
-        create_model=create_model,
-        create_loss=create_loss,
-        create_optimizer=create_optimizer,
-        expected_outputs_dtype=tf.dtypes.float32,
-        data_inputs_dtype=tf.dtypes.float32,
-        print_loss=True
-    ).set_hyperparams(seq2seq_pipeline_hyperparams)
-]).set_name('SignalPrediction')
-
-pipeline = Pipeline([EpochRepeater(
-    ValidationSplitWrapper(
-        MetricsWrapper(Pipeline([
-            TrainOnlyWrapper(DataShuffler()),
-            MiniBatchSequentialPipeline([
-                MetricsWrapper(
-                    signal_prediction_pipeline,
-                    metrics=metrics,
-                    name='batch_metrics'
-                )
-            ], batch_size=batch_size)
-        ]), metrics=metrics,
-            name='epoch_metrics',
-            print_metrics=True
-        ),
-        test_size=validation_size,
-        scoring_function=feature_0_metric
-    ), epochs=epochs)])
+    TrainOnly(DataShuffler()),
+    WindowTimeSeries(),
+    MeanStdNormalizer(),
+    MiniBatchSequentialPipeline([
+        Tensorflow2ModelStep(
+            create_model=create_model,
+            create_loss=create_loss,
+            create_optimizer=create_optimizer,
+            print_loss=True
+        ).set_hyperparams(seq2seq_pipeline_hyperparams)
+    ])
+])
 
 pipeline, outputs = pipeline.fit_transform(data_inputs, expected_outputs)
 ```

neuraxle_tensorflow/tensorflow.py

@@ -3,17 +3,24 @@
 
 class BaseTensorflowModelStep(BaseStep):
     def __init__(
-            self,
-            create_model,
-            create_loss,
-            create_optimizer,
-            step_saver,
-            create_inputs=None,
-            data_inputs_dtype=None,
-            expected_outputs_dtype=None,
-            print_loss=False,
-            print_func=None
+        self,
+        create_model,
+        create_loss,
+        create_optimizer,
+        step_saver,
+        create_inputs=None,
+        data_inputs_dtype=None,
+        expected_outputs_dtype=None,
+        print_loss=False,
+        print_func=None
     ):
+        BaseStep.__init__(
+            self,
+            savers=[step_saver],
+            hyperparams=self.__class__.HYPERPARAMS,
+            hyperparams_space=self.__class__.HYPERPARAMS_SPACE
+        )
+
         self.create_inputs = create_inputs
         self.create_model = create_model
         self.create_loss = create_loss
@@ -22,22 +29,13 @@ def __init__(
         self.expected_outputs_dtype = expected_outputs_dtype
         self.data_inputs_dtype = data_inputs_dtype
 
-        self.set_hyperparams(self.__class__.HYPERPARAMS)
-        self.set_hyperparams_space(self.__class__.HYPERPARAMS_SPACE)
-
         self.train_losses = []
         self.test_losses = []
         self.print_loss = print_loss
         if print_func is None:
             print_func = print
         self.print_func = print_func
 
-        BaseStep.__init__(
-            self,
-            savers=[step_saver],
-            hyperparams=self.HYPERPARAMS
-        )
-
     def add_new_loss(self, loss, test_only=False):
         if test_only:
             if not self.is_train:

neuraxle_tensorflow/tensorflow_v1.py

@@ -69,7 +69,7 @@ def __init__(
         self.has_expected_outputs = has_expected_outputs
         self.create_feed_dict = create_feed_dict
 
-    def setup(self) -> BaseStep:
+    def setup(self, context: ExecutionContext) -> BaseStep:
         """
         Setup tensorflow 1 graph, and session using a variable scope.
 
@@ -84,15 +84,15 @@ def setup(self) -> BaseStep:
             with tf.variable_scope(self.variable_scope, reuse=tf.AUTO_REUSE):
                 self.session = tf.Session(config=tf.ConfigProto(log_device_placement=True), graph=self.graph)
 
-                model = self.create_model(self)
+                model = self.create_model(self, context)
                 if not isinstance(model, tuple):
                     tf.identity(model, name='output')
                 else:
                     tf.identity(model[0], name='output')
                     tf.identity(model[1], name='inference_output')
 
                 tf.identity(self.create_loss(self), name='loss')
-                self.create_optimizer(self).minimize(self['loss'], name='optimizer')
+                self.create_optimizer(self, context).minimize(self['loss'], name='optimizer')
 
                 init = tf.global_variables_initializer()
                 self.session.run(init)
@@ -258,7 +258,7 @@ def load_step(self, step: 'TensorflowV1ModelStep', context: 'ExecutionContext')
         :return: loaded step
         """
         step.is_initialized = False
-        step.setup()
+        step.setup(context)
 
         with step.graph.as_default():
             saver = tf.train.Saver()

neuraxle_tensorflow/tensorflow_v2.py

@@ -70,7 +70,7 @@ def __init__(
             tf_model_checkpoint_folder = 'tensorflow_ckpts'
         self.tf_model_checkpoint_folder = tf_model_checkpoint_folder
 
-    def setup(self) -> BaseStep:
+    def setup(self, context: ExecutionContext) -> BaseStep:
         """
         Setup optimizer, model, and checkpoints for saving.
 
@@ -81,8 +81,8 @@ def setup(self) -> BaseStep:
             return self
 
         with tf.device(self.device_name):
-            self.optimizer = self.create_optimizer(self)
-            self.model = self.create_model(self)
+            self.optimizer = self.create_optimizer(self, context)
+            self.model = self.create_model(self, context)
 
             self.checkpoint = tf.train.Checkpoint(step=tf.Variable(1), optimizer=self.optimizer, net=self.model)
             self.checkpoint_manager = tf.train.CheckpointManager(
@@ -199,7 +199,7 @@ def load_step(self, step: 'Tensorflow2ModelStep', context: 'ExecutionContext') -
         :return: loaded step
         """
         step.is_initialized = False
-        step.setup()
+        step.setup(context)
         step.checkpoint.restore(step.checkpoint_manager.latest_checkpoint)
         return step
 

testing/test_tensorflow_v1.py

@@ -12,7 +12,7 @@
 MATMUL_VARIABLE_SCOPE = "matmul"
 
 
-def create_graph(step: TensorflowV1ModelStep):
+def create_graph(step: TensorflowV1ModelStep, context: ExecutionContext):
     tf.placeholder('float', name='data_inputs')
     tf.placeholder('float', name='expected_outputs')
 
@@ -22,11 +22,11 @@ def create_graph(step: TensorflowV1ModelStep):
     return tf.add(tf.multiply(step['data_inputs'], step['weight']), step['bias'])
 
 
-def create_loss(step: TensorflowV1ModelStep):
+def create_loss(step: TensorflowV1ModelStep, context: ExecutionContext):
     return tf.reduce_sum(tf.pow(step['output'] - step['expected_outputs'], 2)) / (2 * N_SAMPLES)
 
 
-def create_optimizer(step: TensorflowV1ModelStep):
+def create_optimizer(step: TensorflowV1ModelStep, context: ExecutionContext):
     return tf.train.GradientDescentOptimizer(step.hyperparams['learning_rate'])
 
 

testing/test_tensorflow_v2.py

@@ -23,16 +23,16 @@ def toy_dataset():
         dict(x=inputs, y=labels)).repeat(10).batch(2)
 
 
-def create_model(step: Tensorflow2ModelStep):
+def create_model(step: Tensorflow2ModelStep, context: ExecutionContext):
     return LinearModel()
 
 
-def create_optimizer(step: Tensorflow2ModelStep):
+def create_optimizer(step: Tensorflow2ModelStep, context: ExecutionContext):
     return tf.keras.optimizers.Adam(0.1)
 
 
-def create_loss(step: Tensorflow2ModelStep, expected_outputs, actual_outputs):
-    return tf.reduce_mean(tf.abs(actual_outputs - expected_outputs))
+def create_loss(step: Tensorflow2ModelStep, expected_outputs, predicted_outputs):
+    return tf.reduce_mean(tf.abs(predicted_outputs - expected_outputs))
 
 
 def test_tensorflowv2_saver(tmpdir):