Use custom batch dataloader with DDP
Hello,
In our project, we use a custom data loader like so:
class BatchDataLoader(torch.utils.data.dataloader.DataLoader):
    def __init__(self, ds, batch_size, shuffle, num_workers=0):
        inner_sampler = RandomSampler(ds) if shuffle else SequentialSampler(ds)
        sampler = BatchSampler(inner_sampler, batch_size, False)
        super().__init__(ds, batch_size=None, sampler=sampler, num_workers=num_workers)

We use this custom dataloader so that the __getitem__ method is directly called with the list of batch indices and returns a full batch instead of single items. We need this because for some technical reasons, the __getitem__ method is way way faster when accessing a batch of data items rather than single ones. This has been working perfectly fine so far until we started looking into using ddp for training our model.
By default Lightning setup its own sampler for distributed training. We can override this behaviour with replace_sampler_ddp=False but then our custom sampler would be used, whereas it is not made for distributing.
If we wish to make our own custom distributed version of it, then how can we retrieve the distributed context from lightning? Is there a way to know if we are in a distributed context, and if so, how many nodes, etc? Can this be done from a Lightning data module, as this is where we instantiate the data loaders? I tried using torch.distributed,is_available() as a first step. While this passes on Windows, this call fails with AssertionError: Default process group is not initialized on Linux if we are not in a distrbuted context (for instance training on a single GPU). For info, I'm using Lightning 1.0.5 with PyTorch 1.5.1. We wish to setup our sampler automatically, without the user having to specify extra parameters.
Or do you see a way to achieve the same behaviour (getting batches of indices in the __getitem__ method of the dataset) while still benefitting from the default DistributedSampler that Lightning setup? I thought of wrapping my original dataset in another one that would create "buckets" of data, but then how can I make sure the dataset is reshuffled each epoch so that the buckets don't always have the same content?
Thanks for your help :)