Fix the number of training batches used in the training loop (#653)

* Fix the number of processed training batches

* Fix tests

* fix tests

* fix tests

* One more attempt

* Fix another test

pytorch_lightning/trainer/training_loop.py

@@ -367,6 +367,10 @@ class TrainerTrainLoopMixin(ABC):
 
         # run epoch
         for batch_idx, batch in enumerate(self.get_train_dataloader()):
+            # stop epoch if we limited the number of training batches
+            if batch_idx >= self.num_training_batches:
+                break
+
             self.batch_idx = batch_idx
 
             model = self.get_model()
@@ -413,11 +417,6 @@ class TrainerTrainLoopMixin(ABC):
             if early_stop_epoch or self.fast_dev_run:
                 break
 
-            # stop epoch if we limited the number of training batches
-            met_batch_limit = batch_idx >= self.num_training_batches
-            if met_batch_limit:
-                break
-
         # epoch end hook
         if self.is_function_implemented('on_epoch_end'):
             model = self.get_model()

tests/test_logging.py

@@ -309,7 +309,7 @@ def test_custom_logger(tmpdir):
 
     trainer_options = dict(
         max_epochs=1,
-        train_percent_check=0.01,
+        train_percent_check=0.05,
         logger=logger,
         default_save_path=tmpdir
     )

tests/test_restore_models.py

@@ -269,12 +269,12 @@ def test_cpu_restore_training(tmpdir):
     logger = tutils.get_test_tube_logger(tmpdir, False, version=test_logger_version)
 
     trainer_options = dict(
-        max_epochs=2,
+        max_epochs=4,
         val_check_interval=0.50,
         val_percent_check=0.2,
         train_percent_check=0.2,
         logger=logger,
-        checkpoint_callback=ModelCheckpoint(tmpdir)
+        checkpoint_callback=ModelCheckpoint(tmpdir, save_top_k=-1)
     )
 
     # fit model
@@ -290,7 +290,7 @@ def test_cpu_restore_training(tmpdir):
     # we want to see if the weights come back correctly
     new_logger = tutils.get_test_tube_logger(tmpdir, False, version=test_logger_version)
     trainer_options = dict(
-        max_epochs=2,
+        max_epochs=4,
         val_check_interval=0.50,
         val_percent_check=0.2,
         train_percent_check=0.2,

tests/test_trainer.py

@@ -393,7 +393,7 @@ def test_multiple_test_dataloader(tmpdir):
         default_save_path=tmpdir,
         max_epochs=1,
         val_percent_check=0.1,
-        train_percent_check=0.1,
+        train_percent_check=0.2,
     )
 
     # fit model