Clear example for custom DistributedSampler

[Rizhiy]

Hi, I need to train DDP, but I would like to use my own custom DistributedSampler.

I've looked through different issues, but still don't understand how to set up a custom DistributedSampler, I keep getting:
RuntimeError: Default process group has not been initialized, please make sure to call init_process_group.

Can someone please provide a complete example of how to use a custom DistributedSampler?

[carmocca]

Are the PyTorch docs for it not enough?: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler

Note that we do set this sampler for you automatically.

If you could share a snippet of what you are trying to do, it would be easier to help 😄

[Rizhiy]

I have no problem constructing the sampler using PyTorch only, it's only with PL, that I'm having problems.

The sampler you set automatically doesn't work for me, since I need to use custom sampler, and it doesn't appear that you wrap them properly.

Here is a code snippet:

def _create_loader(self, dataset: DictsDataset, batch_size: int, repeat=False, **kwargs) -> DataLoader:
	sampler = Balancedampler(dataset, self.cfg.TRAIN.BALANCE_TARGET)
    if repeat:
    	sampler = RepeatSampler(sampler)
	sampler = DistributedSamplerWrapper(sampler)

    return DataLoader(
    	dataset,
        batch_size=batch_size,
       	num_workers=self._num_workers,
        collate_fn=pred_info_collate,
        worker_init_fn=init_stock_data_worker,
        sampler=sampler,
        **kwargs,
    )

def train_model(self, train_dataset, val_dataset):
	trainer = pl.Trainer(
    	default_root_dir=self.output_dir,
        gpus=self.cfg.NUM_GPUS,
        accelerator="ddp",
        auto_select_gpus=True,
        sync_batchnorm=True,
        max_steps=self.cfg.TRAIN.NUM_ITERS,
        plugins=DDPPlugin(find_unused_parameters=False),
        precision=16 if self.cfg.TRAIN.FP16 else 32,
        val_check_interval=self.cfg.VAL.FREQ,
        replace_sampler_ddp=False,
    )

	trainer.fit(
		self,
        self._create_loader(train_dataset, self.cfg.TRAIN.GPU_BATCH_SIZE, repeat=True),
        self._create_loader(val_dataset, self.cfg.EVAL.GPU_BATCH_SIZE),
    )

[awaelchli]

@Rizhiy Where exactly are you constructing that sampler? The error message tells me you do that too early, at a stage where distributed has not been initialized.
Make sure to construct your sampler inside the train_data_loader methods (either in DataModule or LightningModule).

Hope this answer helps you.

[awaelchli]

# in LightningModule 
def train_dataloader(self):
    dataset = ...
    sampler = DistributedSampler(...)
    dataloader = DataLoader(dataset, sampler=, ...) 
    return dataloader

[Rizhiy]

So all methods defined under DataHooks can be used on LightningModule?

[awaelchli]

Yes, in the link I posted on the right side of the page you have a long list of all hooks that we have.

[Rizhiy]

Ok, thanks.

I think it would a good addition to clarify where DistributedSampler can be added under https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html#replace-sampler-ddp

[awaelchli]

I think that's what @carmocca is doing in his PR here #7636
That's how I found the conversation here <3