Repository: Lightning-AI/pytorch-lightning

Issue:
Add dataloader arg to Trainer.test()
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
It would be nice if you could use a model for inference using:
`Trainer.test(model, test_dataloaders=test_loader)`

### Motivation
This will match the calling structure for `Trainer.fit()` and allow for `test` to be called on any dataset multiple times

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

### Pitch
Here's a use case. After training a model using 5-fold cross-validation, you may want to stack the 5 checkpoints across multiple models, which will require a) out-of-fold (OOF) predictions and b) the 5 test predictions (which will be averaged). It would be cool if a & b could be generated as follows:

```python
for f in folds:
    model1.load_from_checkpoint(f'path/to/model1_fold{f}.ckpt')
    trainer.test(model1,  test_dataloaders=valid_loader)
    trainer.test(model1,  test_dataloaders=test_loader)

    model2.load_from_checkpoint(f'path/to/model2_fold{f}.ckpt'))
    trainer.test(model2,  test_dataloaders=valid_loader)
    trainer.test(model2,  test_dataloaders=test_loader)
```

<!-- A clear and concise description of what you want to happen. -->

### Alternatives
Maybe I'm misunderstanding how `test` works and there is an easier way? Or perhaps the best way to do this is to write an inference function as you would in pure PyTorch?

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
I am in favour of adding this option, but first, lets see how it fits the API
@williamFalcon any strong suggestion against? cc: @PyTorchLightning/core-contributors 
test is meant to ONLY operate on the test set. it’s meant to keep people from using the test set when they shouldn’t haha (ie: only right before publication or right before production use).  

additions that i’m not sure align well
1. Trainer.test as an instance method. Why wouldn’t you just init the trainer? otherwise you won’t be able to test on distributed environments or configure the things you need like apex, etc. 

additions that are good
1. allowing the test function to take in a dataset. this also aligns with how fit works. 
2. fit should also not take a test dataloader (not sure if it does now). 
3. current .test already uses your test dataloader defined in the lightningmodule. so the ONLY addition we’re talking about here is allowing test to ALSO take in a dataloader and use that one only. 
btw I'm interested in how to "train a model using 5-fold cross-validation" in PL.
Let's do this:

1. Add a test_dataloader method to .test()
2. remove the test_dataloader from .fit()?
> btw I'm interested in how to "train a model using 5-fold cross-validation" in PL.

@Ir1d Try this:
https://www.kaggle.com/rohitgr/quest-bert



Resolving patch:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -612,8 +612,7 @@ def fit(
             self,
             model: LightningModule,
             train_dataloader: Optional[DataLoader] = None,
-            val_dataloaders: Optional[DataLoader] = None,
-            test_dataloaders: Optional[DataLoader] = None
+            val_dataloaders: Optional[DataLoader] = None
     ):
         r"""
         Runs the full optimization routine.
@@ -629,14 +628,10 @@ def fit(
                 Pytorch Dataloader or a list of them, specifying validation samples.
                 If the model has a predefined val_dataloaders method this will be skipped
 
-            test_dataloaders: Either a single
-                Pytorch Dataloader or a list of them, specifying validation samples.
-                If the model has a predefined test_dataloaders method this will be skipped
-
         Example::
 
             # Option 1,
-            # Define the train_dataloader(), test_dataloader() and val_dataloader() fxs
+            # Define the train_dataloader() and val_dataloader() fxs
             # in the lightningModule
             # RECOMMENDED FOR MOST RESEARCH AND APPLICATIONS TO MAINTAIN READABILITY
             trainer = Trainer()
@@ -646,15 +641,13 @@ def fit(
             # Option 2
             # in production cases we might want to pass different datasets to the same model
             # Recommended for PRODUCTION SYSTEMS
-            train, val, test = DataLoader(...), DataLoader(...), DataLoader(...)
+            train, val = DataLoader(...), DataLoader(...)
             trainer = Trainer()
             model = LightningModule()
-            trainer.fit(model, train_dataloader=train,
-                        val_dataloader=val, test_dataloader=test)
+            trainer.fit(model, train_dataloader=train, val_dataloader=val)
 
             # Option 1 & 2 can be mixed, for example the training set can be
-            # defined as part of the model, and validation/test can then be
-            # feed to .fit()
+            # defined as part of the model, and validation can then be feed to .fit()
 
         """
         # bind logger and other properties
@@ -662,7 +655,7 @@ def fit(
         self.copy_trainer_model_properties(model)
 
         # set up the passed in dataloaders (if needed)
-        self.__attach_dataloaders(model, train_dataloader, val_dataloaders, test_dataloaders)
+        self.__attach_dataloaders(model, train_dataloader, val_dataloaders)
 
         # check that model is configured correctly
         self.check_model_configuration(model)
@@ -747,7 +740,7 @@ def __set_random_port(self):
             default_port = random.randint(10000, 19000)
             os.environ['MASTER_PORT'] = str(default_port)
 
-    def __attach_dataloaders(self, model, train_dataloader, val_dataloaders, test_dataloaders):
+    def __attach_dataloaders(self, model, train_dataloader=None, val_dataloaders=None, test_dataloaders=None):
         # when dataloader is passed via fit, patch the train_dataloader
         # functions to overwrite with these implementations
         if train_dataloader is not None:
@@ -863,7 +856,7 @@ def run_pretrain_routine(self, model: LightningModule):
         # CORE TRAINING LOOP
         self.train()
 
-    def test(self, model: Optional[LightningModule] = None):
+    def test(self, model: Optional[LightningModule] = None, test_dataloaders: Optional[DataLoader] = None):
         r"""
 
         Separates from fit to make sure you never run on your test set until you want to.
@@ -871,24 +864,36 @@ def test(self, model: Optional[LightningModule] = None):
         Args:
             model: The model to test.
 
+            test_dataloaders: Either a single
+                Pytorch Dataloader or a list of them, specifying validation samples.
+
         Example::
 
             # Option 1
             # run test after fitting
+            test = DataLoader(...)
             trainer = Trainer()
             model = LightningModule()
 
-            trainer.fit()
-            trainer.test()
+            trainer.fit(model)
+            trainer.test(test_dataloaders=test)
 
             # Option 2
             # run test from a loaded model
+            test = DataLoader(...)
             model = LightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')
             trainer = Trainer()
-            trainer.test(model)
+            trainer.test(model, test_dataloaders=test)
         """
 
         self.testing = True
+
+        if test_dataloaders is not None:
+            if model is not None:
+                self.__attach_dataloaders(model, test_dataloaders=test_dataloaders)
+            else:
+                self.__attach_dataloaders(self.model, test_dataloaders=test_dataloaders)
+
         if model is not None:
             self.model = model
             self.fit(model)


Test patch:
diff --git a/tests/trainer/test_dataloaders.py b/tests/trainer/test_dataloaders.py
--- a/tests/trainer/test_dataloaders.py
+++ b/tests/trainer/test_dataloaders.py
@@ -260,12 +260,12 @@ class CurrentTestModel(
     model = CurrentTestModel(hparams)
     trainer = Trainer(**trainer_options)
     fit_options = dict(train_dataloader=model._dataloader(train=True),
-                       val_dataloaders=model._dataloader(train=False),
-                       test_dataloaders=model._dataloader(train=False))
+                       val_dataloaders=model._dataloader(train=False))
+    test_options = dict(test_dataloaders=model._dataloader(train=False))
 
     result = trainer.fit(model, **fit_options)
 
-    trainer.test()
+    trainer.test(**test_options)
 
     assert result == 1
     assert len(trainer.val_dataloaders) == 1, \
@@ -300,11 +300,12 @@ class CurrentTestModel(
     trainer = Trainer(**trainer_options)
     fit_options = dict(train_dataloader=model._dataloader(train=True),
                        val_dataloaders=[model._dataloader(train=False),
-                                        model._dataloader(train=False)],
-                       test_dataloaders=[model._dataloader(train=False),
-                                         model._dataloader(train=False)])
+                                        model._dataloader(train=False)])
+    test_options = dict(test_dataloaders=[model._dataloader(train=False),
+                                          model._dataloader(train=False)])
+
     results = trainer.fit(model, **fit_options)
-    trainer.test()
+    trainer.test(**test_options)
 
     assert len(trainer.val_dataloaders) == 2, \
         f'Multiple `val_dataloaders` not initiated properly, got {trainer.val_dataloaders}'
@@ -342,10 +343,11 @@ class CurrentTestModel(
 
     # fit model
     trainer = Trainer(**trainer_options)
-    fit_options = dict(val_dataloaders=model._dataloader(train=False),
-                       test_dataloaders=model._dataloader(train=False))
+    fit_options = dict(val_dataloaders=model._dataloader(train=False))
+    test_options = dict(test_dataloaders=model._dataloader(train=False))
+
     _ = trainer.fit(model, **fit_options)
-    trainer.test()
+    trainer.test(**test_options)
 
     assert len(trainer.val_dataloaders) == 1, \
         f'`val_dataloaders` not initiated properly, got {trainer.val_dataloaders}'
@@ -511,8 +513,8 @@ class CurrentTestModel(
     )
 
     fit_options = dict(train_dataloader=model._dataloader(train=True),
-                       val_dataloaders=model._dataloader(train=False),
-                       test_dataloaders=model._dataloader(train=False))
+                       val_dataloaders=model._dataloader(train=False))
+    test_options = dict(test_dataloaders=model._dataloader(train=False))
 
     trainer = Trainer(**trainer_options)
 
@@ -524,7 +526,7 @@ class CurrentTestModel(
         trainer.fit(model, **fit_options)
 
     with pytest.warns(UserWarning, match='test'):
-        trainer.test()
+        trainer.test(**test_options)
 
 
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')