Fixed number of epochs model is fitted (Issue #77) (#84)

* Fixed number of epochs model is fitted * Restored nb_epochs to its original value * Changed 'nb_epochs' to 'nb_epoch'
maxpumperla · Apr 19, 2018 · 7562adb · 7562adb
1 parent 31b1f91
commit 7562adb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/elephas/spark_model.py b/elephas/spark_model.py
@@ -304,8 +304,10 @@ def train(self, data_iterator):
                 weights_before_training = get_server_weights(self.master_url)
                 model.set_weights(weights_before_training)
                 self.train_config['epochs'] = 1
+                self.train_config['nb_epoch'] = 1
                 if x_train.shape[0] > batch_size:
                     model.fit(x_train, y_train, **self.train_config)
+                self.train_config['nb_epoch'] = nb_epoch
                 weights_after_training = model.get_weights()
                 deltas = subtract_params(weights_before_training, weights_after_training)
                 put_deltas_to_server(deltas, self.master_url)