Repository: Lightning-AI/pytorch-lightning

Issue:
0.7.3 breaks reusable dataloaders in DDP
## 🐛 Bug

0.7.3 breaks reusable dataloaders in DDP

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 345, in ddp_train
    self.run_pretrain_routine(model)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 864, in run_pretrain_routine
    self.train()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 296, in train
    self.reset_train_dataloader(model)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/data_loading.py", line 128, in reset_train_dataloader
    self.train_dataloader = self.auto_add_sampler(self.train_dataloader, train=True)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/data_loading.py", line 112, in auto_add_sampler
    dataloader = type(dataloader)(**dl_args)
  File "../main/dataset.py", line 15, in __init__
    super().__init__(*args, **kwargs)
TypeError: __init__() got an unexpected keyword argument 'iterator'
```

#### Code sample

```
class _RepeatSampler(object):
    def __init__(self, sampler):
        self.sampler = sampler

    def __iter__(self):
        while True:
            yield from iter(self.sampler)

class FastDataLoader(torch.utils.data.dataloader.DataLoader):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
        self.iterator = super().__iter__()

    def __len__(self):
        return len(self.batch_sampler.sampler)

    def __iter__(self):
        for i in range(len(self)):
            yield next(self.iterator)
```
replace Dataloader with FastDataLoader in lightning
(this snippet is from https://github.com/pytorch/pytorch/issues/15849)

### Expected behavior

Dataloaders initialize correctly and are reused between train/val/epochs (works as expected in 0.7.1)

### Probable Cause

https://github.com/PyTorchLightning/pytorch-lightning/pull/1425


Discussion:
ummm yeah. we should change the dataloader swap with swapping a dataloader init from the class or not swipe the dataloder at all but set the correct sampler. 

@justusschock any ideas?
This is a mixture of  #1425 and #1346 

And I don't think we can prevent this when we want to set correct samplers also in subclasses of `DataLoader`. We use all public attributes for reinitialization.

The probably easiest fix for you, would be to change `self.iterator` to `self._iterator` to avoid passing this argument in reinit.

If we just change the sampler, this might yield unexpected behaviour.

Resolving patch:
diff --git a/pytorch_lightning/trainer/data_loading.py b/pytorch_lightning/trainer/data_loading.py
--- a/pytorch_lightning/trainer/data_loading.py
+++ b/pytorch_lightning/trainer/data_loading.py
@@ -61,6 +61,7 @@ class TrainerDataLoadingMixin(ABC):
     train_percent_check: float
     val_percent_check: float
     test_percent_check: float
+    replace_sampler_ddp: bool
 
     @abstractmethod
     def is_overriden(self, *args):
@@ -88,10 +89,8 @@ def auto_add_sampler(self, dataloader: DataLoader, train: bool) -> DataLoader:
         # don't do anything if it's not a dataloader
         if not isinstance(dataloader, DataLoader):
             return dataloader
-
-        need_dist_sampler = self.use_ddp or self.use_ddp2 or self.use_tpu
-
-        if need_dist_sampler:
+        need_dist_sampler = (self.use_ddp or self.use_ddp2 or self.use_tpu)
+        if self.replace_sampler_ddp and need_dist_sampler:
 
             skip_keys = ['sampler', 'batch_sampler', 'dataset_kind']
 
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
@@ -127,6 +127,7 @@ def __init__(
             benchmark: bool = False,
             reload_dataloaders_every_epoch: bool = False,
             auto_lr_find: Union[bool, str] = False,
+            replace_sampler_ddp: bool = True,
             default_save_path=None,  # backward compatible, todo: remove in v0.8.0
             gradient_clip=None,  # backward compatible, todo: remove in v0.8.0
             nb_gpu_nodes=None,  # backward compatible, todo: remove in v0.8.0
@@ -282,6 +283,9 @@ def __init__(
                 rate in self.hparams.lr | self.hparams.learning_rate in the lightning module.
                 To use a different key, set a string instead of True with the key name.
 
+            replace_sampler_ddp: Explicitly enables or disables sampler replacement.
+                If not specified this will toggled automatically ddp is used
+
             benchmark: If true enables cudnn.benchmark.
 
             terminate_on_nan: If set to True, will terminate training (by raising a `ValueError`) at the
@@ -362,6 +366,7 @@ def __init__(
         self.reload_dataloaders_every_epoch = reload_dataloaders_every_epoch
 
         self.auto_lr_find = auto_lr_find
+        self.replace_sampler_ddp = replace_sampler_ddp
 
         self.truncated_bptt_steps = truncated_bptt_steps
         self.resume_from_checkpoint = resume_from_checkpoint


Test patch: