Pytorch Lightning DDP Sampling - How to correctly do it with torchgeo

[TolgaAktas]

I am working on migrating my codebase to PLightning for easier experimentation and multi-GPU training, but I learned that PL with DDP backend requires some gizmos with the sampler. I believe that I need to provide a custom sampler to DDP trainer, which I believe would be one of the GeoSamplers of torchgeo. Is that so? I have the following sampler and dataloader in my vanilla PyTorch script, and I was wondering if I could just pass those to the PyTorch Trainer for DDP.

train_dataset = IntersectionDataset(shadow_train_set,clean_train_set,transforms=transforms)
train_sampler = RandomBatchGeoSampler(train_dataset,
                                      size=(opt.train_ps,opt.train_ps),
                                      batch_size=opt.batch_size,
                                      length=opt.num_batches * opt.batch_size)

train_loader = DataLoader(train_dataset, 
                          batch_sampler = train_sampler, 
                          collate_fn=pick_images,
                          num_workers=opt.train_workers,
                          pin_memory=True,
                          drop_last=False)

shadow_valid_set = Landsat8(root=opt.shadow_valid_dir,bands=bands)
clean_valid_set  = Landsat8(root=opt.clean_valid_dir,bands=bands)

# val_dataset = shadow_valid_set & clean_valid_set
val_dataset = IntersectionDataset(shadow_valid_set,clean_valid_set,transforms=transforms)
val_sampler = RandomBatchGeoSampler(val_dataset,size=(opt.train_ps,opt.train_ps),batch_size=1,length=opt.num_batches)
val_loader = DataLoader(val_dataset, 
                        batch_sampler = val_sampler, 
                        collate_fn=pick_images,
                        num_workers=opt.eval_workers,
                        pin_memory=False,
                        drop_last=False)

[isaaccorley]

Our samplers don't currently support distributed training but this is on our roadmap.

[TolgaAktas]

So I wouldn't be able to do DDP for single-machine multi-GPU training? Or is there a hacky way to solve this?

[isaaccorley]

I haven't tried in awhile but I believe PL actually now wraps your sampler as a distributed sampler automatically if you use DDP so maybe just try it out and see if it works!

[adamjstewart]

Related to #305

[TolgaAktas]

So I should be fine with just doing this?

train_dataset = IntersectionDataset(shadow_train_set,clean_train_set,transforms=transforms)
train_sampler = RandomBatchGeoSampler(train_dataset,
                                      size=(opt.train_ps,opt.train_ps),
                                      batch_size=opt.batch_size,
                                      length=opt.num_batches * opt.batch_size)

train_loader = DataLoader(train_dataset, 
                          batch_sampler = train_sampler, 
                          collate_fn=pick_images,
                          num_workers=opt.train_workers,
                          pin_memory=True,
                          drop_last=False)

shadow_valid_set = Landsat8(root=opt.shadow_valid_dir,bands=bands)
clean_valid_set  = Landsat8(root=opt.clean_valid_dir,bands=bands)

# val_dataset = shadow_valid_set & clean_valid_set
val_dataset = IntersectionDataset(shadow_valid_set,clean_valid_set,transforms=transforms)
val_sampler = RandomBatchGeoSampler(val_dataset,size=(opt.train_ps,opt.train_ps),batch_size=1,length=opt.num_batches)
val_loader = DataLoader(val_dataset, 
                        batch_sampler = val_sampler, 
                        collate_fn=pick_images,
                        num_workers=opt.eval_workers,
                        pin_memory=False,
                        drop_last=False)

trainer = Trainer(gpus=-1,distributed_backend="ddp",......)
trainer.fit(model,train_dataloaders=train_loader,val_dataloaders=val_loader,....)

[adamjstewart]

Try it and see.