generalize reinstantiation of dataloader (Lightning-AI#1346)

* generalize reinstantiation of dataloader

* fix condition

* add test

* update changelog

* fix changelog

Co-authored-by: J. Borovec <[email protected]>

Author: 2 people

pytorch_lightning/trainer/data_loading.py

@@ -84,16 +84,10 @@ def auto_add_sampler(self, dataloader: DataLoader, train: bool) -> DataLoader:
 
         if need_dist_sampler and no_sampler_added:
 
+            skip_keys = ['sampler', 'batch_sampler', 'dataset_kind']
+
             dl_args = {
-                'dataset': dataloader.dataset,
-                'batch_size': dataloader.batch_size,
-                'shuffle': False,
-                'num_workers': dataloader.num_workers,
-                'collate_fn': dataloader.collate_fn,
-                'pin_memory': dataloader.pin_memory,
-                'drop_last': dataloader.drop_last,
-                'timeout': dataloader.timeout,
-                'worker_init_fn': dataloader.worker_init_fn
+                k: v for k, v in dataloader.__dict__.items() if not k.startswith('_') and k not in skip_keys
             }
 
             if self.use_tpu:
@@ -102,13 +96,11 @@ def auto_add_sampler(self, dataloader: DataLoader, train: bool) -> DataLoader:
                     num_replicas=xm.xrt_world_size(),
                     rank=xm.get_ordinal()
                 )
-                dl_args['shuffle'] = False
             else:
                 sampler = DistributedSampler(dataloader.dataset)
-                dl_args['shuffle'] = False
 
             dl_args['sampler'] = sampler
-            dataloader = DataLoader(**dl_args)
+            dataloader = type(dataloader)(**dl_args)
 
         return dataloader
 

tests/trainer/test_dataloaders.py

@@ -1,4 +1,5 @@
 import pytest
+import torch
 
 import tests.base.utils as tutils
 from pytorch_lightning import Trainer
@@ -482,3 +483,41 @@ class CurrentTestModel(
             test_percent_check=0.5
         )
         trainer.fit(model)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')
+def test_dataloader_reinit_for_subclass():
+
+    class CustomDataLoader(torch.utils.data.DataLoader):
+        def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,
+                     batch_sampler=None, num_workers=0, collate_fn=None,
+                     pin_memory=False, drop_last=False, timeout=0,
+                     worker_init_fn=None, dummy_kwarg=None):
+            super().__init__(dataset,
+                             batch_size,
+                             shuffle,
+                             sampler,
+                             batch_sampler,
+                             num_workers,
+                             collate_fn,
+                             pin_memory,
+                             drop_last,
+                             timeout,
+                             worker_init_fn)
+
+            self.dummy_kwarg = dummy_kwarg
+
+    trainer = Trainer(gpus=[0, 1],
+                      num_nodes=1,
+                      distributed_backend='ddp')
+
+    class CustomDummyObj:
+        sampler = None
+
+    result = trainer.auto_add_sampler(CustomDummyObj(), train=True)
+    assert isinstance(result, CustomDummyObj), "Wrongly reinstantiated data loader"
+
+    result = trainer.auto_add_sampler(CustomDataLoader(list(range(1000))), train=True)
+    assert isinstance(result, torch.utils.data.DataLoader)
+    assert isinstance(result, CustomDataLoader)
+    assert hasattr(result, 'dummy_kwarg')