Add NeuralNet.partial_fit method that trains for only one epoch

nolearn/lasagne/base.py

@@ -403,7 +403,7 @@ def _create_iter_funcs(self, layers, objective, update, output_type):
 
         return train_iter, eval_iter, predict_iter
 
-    def fit(self, X, y):
+    def fit(self, X, y, epochs=None):
         X, y = self._check_good_input(X, y)
 
         if self.use_label_encoder:
@@ -413,12 +413,16 @@ def fit(self, X, y):
         self.initialize()
 
         try:
-            self.train_loop(X, y)
+            self.train_loop(X, y, epochs=epochs)
         except KeyboardInterrupt:
             pass
         return self
 
-    def train_loop(self, X, y):
+    def partial_fit(self, X, y, classes=None):
+        return self.fit(X, y, epochs=1)
+
+    def train_loop(self, X, y, epochs=None):
+        epochs = epochs or self.max_epochs
         X_train, X_valid, y_train, y_valid = self.train_split(X, y, self)
 
         on_epoch_finished = self.on_epoch_finished
@@ -447,7 +451,7 @@ def train_loop(self, X, y):
 
         num_epochs_past = len(self.train_history_)
 
-        while epoch < self.max_epochs:
+        while epoch < epochs:
             epoch += 1
 
             train_losses = []

nolearn/lasagne/tests/test_base.py

@@ -105,6 +105,14 @@ def test_load_params_from_message(self, net, net_fitted, capsys):
 """
         assert out == message
 
+    def test_partial_fit(self, net, X_train, y_train):
+        net2 = clone(net)
+        assert net2.partial_fit(X_train, y_train) is net2
+        net2.partial_fit(X_train, y_train)
+        history = net2.train_history_
+        assert len(history) == 2
+        assert history[1]['valid_accuracy'] > 0.85
+
 
 def test_lasagne_functional_grid_search(mnist, monkeypatch):
     # Make sure that we can satisfy the grid search interface.