release v0.1.dev16

pytorch_lightning/models/trainer.py

@@ -270,21 +270,22 @@ class Trainer(TrainerIO):
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
-                self.__run_tng_batch(data_batch)
+                batch_result = self.__run_tng_batch(data_batch)
+                early_stop_epoch = batch_result == -1
 
                 # ---------------
                 # RUN VAL STEP
                 # ---------------
                 is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
-                if self.fast_dev_run or is_val_check_batch:
+                if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
                     self.__run_validation()
 
                 # when batch should be saved
-                if (batch_nb + 1) % self.log_save_interval == 0:
+                if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
                     self.experiment.save()
 
                 # when metrics should be logged
-                if batch_nb % self.add_log_row_interval == 0:
+                if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
@@ -308,6 +309,10 @@ class Trainer(TrainerIO):
                 if self.__is_function_implemented('on_batch_end'):
                     self.model.on_batch_end()
 
+                # end epoch early
+                if early_stop_epoch:
+                    break
+
             # hook
             if self.__is_function_implemented('on_epoch_end'):
                 self.model.on_epoch_end()
@@ -322,15 +327,16 @@ class Trainer(TrainerIO):
                 if stop:
                     return
 
+
     def __run_tng_batch(self, data_batch):
         if data_batch is None:
-            return
+            return 0
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
             response = self.model.on_batch_start(data_batch)
             if response == -1:
-                return
+                return -1
 
         if self.enable_tqdm:
             self.prog_bar.update(1)
@@ -372,6 +378,8 @@ class Trainer(TrainerIO):
         if self.__is_function_implemented('on_batch_end'):
             self.model.on_batch_end()
 
+        return 0
+
     def __run_validation(self):
         # decide if can check epochs
         can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0

pytorch_lightning/root_module/model_saving.py

@@ -88,6 +88,7 @@ class TrainerIO(object):
         self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
         self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
         self.global_step = checkpoint['global_step']
+        self.current_epoch = checkpoint['epoch']
 
         # restore the optimizers
         optimizer_states = checkpoint['optimizer_states']

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.1.dev14',
+    version='0.1.dev16',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",