Add missing test for "multiple dataloader + percent_check fix" (Lightning-AI#2226)

* Init fix num_batches

* Fix num_batches in case of multiple dataloaders

* Apply suggestions from code review

* Changes based on suggestions

* Flake8

* Add test to check num_batches

* generalize dataloader percent check test

* fix formatting

* remove hparams

* tests

* CHANGELOG

* Update CHANGELOG.md

* max_batches can be int

* conflict and rebase

* add back the test


fix


fix message


0.0 works


Revert "fix message"

This reverts commit 839cacf8b8610f4e697e654ef6f3d2501bf23984.

* update changelog

* Update CHANGELOG.md

* Fix num batches in case of multiple dataloaders and percent_check (Lightning-AI#1920)

* git conflict

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>

* missing union

* doc update suggestion by @rohitgr7

* extend test

* changelog

* docs add note about multiple loaders

* update changelog

* remove unused variable

Co-authored-by: rohitgr7 <rohitgr1998@gmail.com>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -18,6 +18,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed number batches in case of multiple dataloaders and `limit_{*}_batches` ([#1920](https://github.com/PyTorchLightning/pytorch-lightning/pull/1920), [#2226](https://github.com/PyTorchLightning/pytorch-lightning/pull/2226))
+
 - Fixed an issue with forward hooks not being removed after model summary ([#2298](https://github.com/PyTorchLightning/pytorch-lightning/pull/2298))
 
 - Fixed ROC metric for CUDA tensors ([#2304](https://github.com/PyTorchLightning/pytorch-lightning/pull/2304))

pytorch_lightning/trainer/__init__.py

@@ -456,6 +456,8 @@ def on_train_end(self, trainer, pl_module):
     # run for only 10 batches
     trainer = Trainer(limit_test_batches=10)
 
+In the case of multiple test dataloaders, the limit applies to each dataloader individually.
+
 limit_val_batches
 ^^^^^^^^^^^^^^^^^
 
@@ -473,6 +475,8 @@ def on_train_end(self, trainer, pl_module):
     # run for only 10 batches
     trainer = Trainer(limit_val_batches=10)
 
+In the case of multiple validation dataloaders, the limit applies to each dataloader individually.
+
 log_gpu_memory
 ^^^^^^^^^^^^^^
 Options:

pytorch_lightning/trainer/data_loading.py

@@ -287,8 +287,6 @@ def _reset_eval_dataloader(
             for i, dataloader in enumerate(dataloaders):
                 num_batches = 0
                 self._worker_check(dataloader, f'{mode} dataloader {i}')
-                if not _has_len(dataloader):
-                    num_batches = float('inf')
 
                 # percent or num_steps
                 limit_eval_batches = getattr(self, f'limit_{mode}_batches')

pytorch_lightning/trainer/evaluation_loop.py

@@ -124,7 +124,7 @@
 
 from abc import ABC, abstractmethod
 from pprint import pprint
-from typing import Callable, Optional, List
+from typing import Callable, Optional, List, Union
 
 import torch
 from torch.utils.data import DataLoader
@@ -222,13 +222,20 @@ def reset_test_dataloader(self, *args):
     def reset_val_dataloader(self, *args):
         """Warning: this is just empty shell for code implemented in other class."""
 
-    def _evaluate(self, model: LightningModule, dataloaders, max_batches: List[int], test_mode: bool = False):
+    def _evaluate(
+        self,
+        model: LightningModule,
+        dataloaders: List[DataLoader],
+        max_batches: Union[int, List[int]],
+        test_mode: bool = False
+    ):
         """Run evaluation code.
 
         Args:
-            model: PT model
-            dataloaders: list of PT dataloaders
-            max_batches: List of scalars
+            model: The model to evaluate.
+            dataloaders: A list of PyTorch dataloaders.
+            max_batches: An integer or list of integers with length of the number of dataloaders. Each
+                entry is the number of batches to process in the corresponding dataloader.
             test_mode:
         """
         # enable eval mode
@@ -244,6 +251,10 @@ def _evaluate(self, model: LightningModule, dataloaders, max_batches: List[int],
         # bookkeeping
         outputs = []
 
+        # convert max_batches to list
+        if isinstance(max_batches, int):
+            max_batches = [max_batches] * len(dataloaders)
+
         # run validation
         for dataloader_idx, dataloader in enumerate(dataloaders):
             dl_outputs = []

pytorch_lightning/trainer/trainer.py

@@ -223,7 +223,7 @@ def __init__(
 
             min_steps: Force training for at least these number of steps. Disabled by default (None).
 
-            limit_train_batches: How much of training dataset to check.
+            limit_train_batches: How much of training dataset to check (floats = percent, int = num_batches)
 
             limit_val_batches: How much of validation dataset to check (floats = percent, int = num_batches)
 

pytorch_lightning/trainer/training_loop.py

@@ -208,7 +208,6 @@ class TrainerTrainLoopMixin(ABC):
     check_val_every_n_epoch: ...
     num_training_batches: int
     val_check_batch: ...
-    num_val_batches: int
     disable_validation: bool
     fast_dev_run: ...
     accumulation_scheduler: ...

tests/base/model_test_dataloaders.py

@@ -7,7 +7,7 @@
 class TestDataloaderVariations(ABC):
 
     @abstractmethod
-    def dataloader(self, train: bool):
+    def dataloader(self, *args, **kwargs):
         """placeholder"""
 
     def test_dataloader(self):
@@ -19,6 +19,11 @@ def test_dataloader__infinite(self):
     def test_dataloader__not_implemented_error(self):
         return CustomNotImplementedErrorDataloader(self.dataloader(train=False))
 
+    def test_dataloader__multiple_mixed_length(self):
+        lengths = [50, 30, 40]
+        dataloaders = [self.dataloader(train=False, num_samples=n) for n in lengths]
+        return dataloaders
+
     def test_dataloader__empty(self):
         return None
 

tests/base/model_test_epoch_ends.py

@@ -7,7 +7,7 @@ class TestEpochEndVariations(ABC):
 
     def test_epoch_end(self, outputs):
         """
-        Called at the end of validation to aggregate outputs
+        Called at the end of test epoch to aggregate outputs
         :param outputs: list of individual outputs of each validation step
         :return:
         """
@@ -40,7 +40,7 @@ def test_epoch_end(self, outputs):
 
     def test_epoch_end__multiple_dataloaders(self, outputs):
         """
-        Called at the end of validation to aggregate outputs
+        Called at the end of test epoch to aggregate outputs
         :param outputs: list of individual outputs of each validation step
         :return:
         """

tests/base/model_utilities.py

@@ -6,8 +6,8 @@
 class ModelTemplateData:
     hparams: ...
 
-    def dataloader(self, train):
-        dataset = TrialMNIST(root=self.data_root, train=train, download=True)
+    def dataloader(self, train: bool, num_samples: int = 100):
+        dataset = TrialMNIST(root=self.data_root, train=train, num_samples=num_samples, download=True)
 
         loader = DataLoader(
             dataset=dataset,

tests/base/model_valid_dataloaders.py

@@ -7,12 +7,17 @@
 class ValDataloaderVariations(ABC):
 
     @abstractmethod
-    def dataloader(self, train: bool):
+    def dataloader(self, *args, **kwargs):
         """placeholder"""
 
     def val_dataloader(self):
         return self.dataloader(train=False)
 
+    def val_dataloader__multiple_mixed_length(self):
+        lengths = [100, 30]
+        dataloaders = [self.dataloader(train=False, num_samples=n) for n in lengths]
+        return dataloaders
+
     def val_dataloader__multiple(self):
         return [self.dataloader(train=False),
                 self.dataloader(train=False)]