Repository: Lightning-AI/pytorch-lightning

Issue:
num_training_batches rounds down, causing 0 batches count
## 🐛 Bug

self.num_training_batches is defined using int [here](https://github.com/williamFalcon/pytorch-lightning/blob/ca73b70d15bc8db3f57c1fd2d3bf152e6e1d7c4e/pytorch_lightning/trainer/data_loading.py#L52), which rounds it down to 0 when a small training_percent_check or overfit_pct is used, even though at least 1 batch is still processed. 

This does not cause any errors in "vanilla" lightning, but crashes any user code that uses the number of batches in a division (for example to get an average of some quantity over batches).

### To Reproduce

Steps to reproduce the behavior:

Set the training percentage to a small enough percentage that the number of examples is smaller than the batch size for a given dataset.

This would require a very simple fix, either to use `math.ceil()` or `max(1, self.num_training_batches)`, depending of how the quantity is expected to behave in the rest of the code.



Discussion:
I think `max(1, self.num_training_batches)` would be the best because it doesn't change the behaviour for existing code that has num_training_batches >= 1. I'm sure the PL team would appreciate if you made a PR. Is the bug also present for validation and test batches?
Well, actually it is not the rounding issue. The problem is that now we process `num_training_batches + 1` batches instead of `num_training_batches` if `num_training_batches < len(train_dataloader)`.

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -367,6 +367,10 @@ def run_training_epoch(self):
 
         # 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 @@ def run_training_epoch(self):
             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()


Test patch:
diff --git a/tests/test_logging.py b/tests/test_logging.py
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -309,7 +309,7 @@ def version(self):
 
     trainer_options = dict(
         max_epochs=1,
-        train_percent_check=0.01,
+        train_percent_check=0.05,
         logger=logger,
         default_save_path=tmpdir
     )
diff --git a/tests/test_restore_models.py b/tests/test_restore_models.py
--- a/tests/test_restore_models.py
+++ b/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,
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -393,7 +393,7 @@ class CurrentTestModel(
         default_save_path=tmpdir,
         max_epochs=1,
         val_percent_check=0.1,
-        train_percent_check=0.1,
+        train_percent_check=0.2,
     )
 
     # fit model