fix batchsize when using BatchSampler

Author: gegejun

src/lightning/pytorch/utilities/data.py

@@ -169,9 +169,9 @@ def _get_dataloader_init_args_and_kwargs(
         if was_wrapped:
             # if the dataloader was wrapped in a hook, only take arguments with default values
             # and assume user passes their kwargs correctly
-            params.update({
-                k: v for k, v in inspect.signature(DataLoader.__init__).parameters.items() if v.default is not v.empty
-            })
+            params.update(
+                {k: v for k, v in inspect.signature(DataLoader.__init__).parameters.items() if v.default is not v.empty}
+            )
         else:
             params.update(inspect.signature(DataLoader.__init__).parameters)
             params.pop("self", None)
@@ -332,6 +332,17 @@ def _dataloader_init_kwargs_resolve_sampler(
             "batch_size": 1,
             "drop_last": False,
         }
+    if batch_sampler is not None and batch_sampler_cls is BatchSampler:
+        # This is a PyTorch `BatchSampler` but maybe created by user, so batch_size and drop_last should be preserved
+        batch_size = batch_sampler.batch_size
+        drop_last = batch_sampler.drop_last if not is_predicting else False
+        return {
+            "sampler": sampler,
+            "shuffle": False,
+            "batch_sampler": None,
+            "batch_size": batch_size,
+            "drop_last": drop_last,
+        }
 
     return {"sampler": sampler, "shuffle": False, "batch_sampler": None}
 

tests/tests_pytorch/trainer/test_batch_sampler.py

@@ -0,0 +1,39 @@
+from torch.utils.data import RandomSampler, BatchSampler
+from torch.utils.data.dataloader import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from lightning.pytorch import Callback, Trainer, seed_everything
+from tests_pytorch.helpers.runif import RunIf
+from lightning.pytorch.demos.boring_classes import (
+    BoringModel,
+    RandomDataset,
+)
+
+
+class DistribBatchSamplerCallback(Callback):
+    def __init__(self, expected_batch_size=4):
+        self.expected_batch_size = expected_batch_size
+
+    def on_train_start(self, trainer, pl_module):
+        assert isinstance(trainer.train_dataloader.batch_sampler, DistributedSampler)
+        assert trainer.train_dataloader.batch_size == self.expected_batch_size
+
+
+@RunIf(min_cuda_gpus=2, skip_windows=True)
+def test_dataloader_distributed_batch_sampler(tmp_path):
+    """Test DistributedSampler and it's arguments for DDP backend."""
+    seed_everything(123)
+    dataset = RandomDataset(32, 64)
+    sampler = RandomSampler(dataset)
+    batch_sampler = BatchSampler(sampler, batch_size=4, drop_last=False)
+    dataloader = DataLoader(dataset, batch_sampler=batch_sampler)
+    model = BoringModel()
+    trainer = Trainer(
+        accelerator="gpu",
+        devices=[0, 1],
+        num_nodes=1,
+        strategy="ddp",
+        default_root_dir=tmp_path,
+        max_steps=1,
+        callbacks=[DistribBatchSamplerCallback(expected_batch_size=4)],
+    )
+    trainer.fit(model, train_dataloaders=dataloader)