Repository: Lightning-AI/pytorch-lightning

Issue:
Only invoke setup() once, not in both trainer.fit() and trainer.test()
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
Only invoke `def setup(self, step: str)` when calling `trainer.test(net)` if this has not been called before (e.g. `trainer.fit(net)`).

### Motivation

<!-- 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 -->
The `setup` function is [described in the docs](https://pytorch-lightning.readthedocs.io/en/latest/introduction_guide.html#models-defined-by-data) as: "use setup to do splits, and build your model internals".
Therefore, I wrote code that does the train-val-test function and some DataFrame labels transformation (e.g. label to one_hot) in this function.
A pretty common pattern is the following:
```python
trainer.fit(net)
trainer.test(net)
```
Contrary to what I expected, I saw from my debug output that `setup` was invoked twice.
This is a **waste of computational resources**, and since I did the train-val split randomly, I do not have access to the indices that were used in either step (and possibly other issues such as the label transformation re-ordering the columns of which number represent which label).

Current situation, train and val step use same `setup`, test step uses another invoke of `setup`
I assume that it is more common that the train-val and test step of the trainer use the same setup code, than that setup does something special for only test (and not val).
As Lightning works by giving sensible defaults, and allowing you to hack at anything you want, the logic should be that `setup` should only be invoked once, and allow for a way to specify a special test `setup` function.

### Pitch

<!-- A clear and concise description of what you want to happen. -->
Have the trainer keep track of whether `setup()` has been invoked or not, so `setup()` can be skipped in `trainer.test(net)` if it was already invoked in `trainer.fit(net)`. This way the common use case will benefit from less computation power and is more in line what is expected of the magic of Lightning.

I'm not sure what would be the best approach for users to set a separate setup call for testing.
Maybe something like: `trainer.test(net, always_invoke_setup=True)`?

### Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
Include some custom logic that checks if data has been initialized:
```python
def setup(self, step: str):
    if self.data is None:
        # setup code
    else:
        # do nothing
```

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->
Ideas formed by discussing this issue on the pytorch-lightning SLACK in the questions channel. Thanks goed to the people who replied.


Discussion:
we could also just split the method

fit_setup
test_setup?

@tullie 
Or simply just don't calll `setup('fit')` when `self.testing == True`??
that doesn’t happen no?

it you call .test() only the setup(‘test’) gets called 
Since it calls `.fit()` within test, it calls `.setup('fit')` too, I think:
https://github.com/PyTorchLightning/pytorch-lightning/blob/7b4db3045dcc9e6bb0b66e409b25bb2c7fa378f0/pytorch_lightning/trainer/trainer.py#L1033-L1048
oh! yeah that's a bug :) 

mind submitting a PR? 

nice catch!!

Resolving patch:
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
--- a/pytorch_lightning/trainer/distrib_data_parallel.py
+++ b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -506,8 +506,8 @@ def ddp_train(self, process_idx, q, model, is_master=False, proc_offset=0):
         model.init_ddp_connection(self.global_rank, self.world_size, self.is_slurm_managing_tasks)
 
         # call setup after the ddp process has connected
-        self.setup('fit')
-        if self.is_function_implemented('setup', model):
+        if not self.testing:
+            self.setup('fit')
             model.setup('fit')
 
         # on world_size=0 let everyone know training is starting
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
--- a/pytorch_lightning/trainer/distrib_parts.py
+++ b/pytorch_lightning/trainer/distrib_parts.py
@@ -167,8 +167,8 @@ def __transfer_batch_to_device(self, batch: Any, device: torch.device):
 
     def single_gpu_train(self, model):
         # call setup
-        self.setup('fit')
-        if self.is_function_implemented('setup', model):
+        if not self.testing:
+            self.setup('fit')
             model.setup('fit')
 
         model.cuda(self.root_gpu)
@@ -189,8 +189,8 @@ def single_gpu_train(self, model):
 
     def tpu_train(self, tpu_core_idx, model):
         # call setup after the ddp process has connected
-        self.setup('fit')
-        if self.is_function_implemented('setup', model):
+        if not self.testing:
+            self.setup('fit')
             model.setup('fit')
 
         # put model on tpu
@@ -229,8 +229,8 @@ def tpu_train(self, tpu_core_idx, model):
 
     def dp_train(self, model):
         # call setup after the ddp process has connected
-        self.setup('fit')
-        if self.is_function_implemented('setup', model):
+        if not self.testing:
+            self.setup('fit')
             model.setup('fit')
 
         model.cuda(self.root_gpu)
@@ -275,8 +275,8 @@ def dp_train(self, model):
 
     def horovod_train(self, model):
         # call setup after the ddp process has connected
-        self.setup('fit')
-        if self.is_function_implemented('setup', model):
+        if not self.testing:
+            self.setup('fit')
             model.setup('fit')
 
         if torch.cuda.is_available() and self.on_gpu:
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
@@ -1054,8 +1054,8 @@ def fit(
                 raise MisconfigurationException('amp + cpu is not supported.  Please use a GPU option')
 
             # call setup after the ddp process has connected
-            self.setup('fit')
-            if self.is_function_implemented('setup', model):
+            if not self.testing:
+                self.setup('fit')
                 model.setup('fit')
 
             # CHOOSE OPTIMIZER
@@ -1328,8 +1328,7 @@ def test(
 
     def __test_using_best_weights(self, ckpt_path, test_dataloaders):
         model = self.get_model()
-        if self.is_function_implemented('setup', model):
-            model.setup('test')
+        model.setup('test')
 
         # if user requests the best checkpoint but we don't have it, error
         if ckpt_path == 'best' and self.checkpoint_callback.save_top_k <= 0:
@@ -1373,8 +1372,7 @@ def __test_using_best_weights(self, ckpt_path, test_dataloaders):
 
     def __test_given_model(self, model, test_dataloaders):
         # setup hook
-        if self.is_function_implemented('setup', model):
-            model.setup('test')
+        model.setup('test')
 
         # attach data
         if test_dataloaders is not None:


Test patch:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -980,3 +980,34 @@ def test_trainer_pickle(tmpdir):
     )
     pickle.dumps(trainer)
     cloudpickle.dumps(trainer)
+
+
+def test_trainer_setup_call(tmpdir):
+    """Test setup call with fit and test call."""
+
+    class CurrentModel(EvalModelTemplate):
+
+        def setup(self, stage):
+            self.stage = stage
+
+    class TrainerSubclass(Trainer):
+
+        def setup(self, stage):
+            self.stage = stage
+
+    model = CurrentModel()
+
+    # fit model
+    trainer = TrainerSubclass(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        checkpoint_callback=False
+    )
+
+    trainer.fit(model)
+    assert trainer.stage == 'fit'
+    assert trainer.get_model().stage == 'fit'
+
+    trainer.test(ckpt_path=None)
+    assert trainer.stage == 'test'
+    assert trainer.get_model().stage == 'test'