Lightning-AI/pytorch-lightning

torch.randn() + DDP + GANs are easy to get wrong with lightning

Open

#16.166 aperta il 21 dic 2022

Vedi su GitHub
 (3 commenti) (2 reazioni) (0 assegnatari)Python (3233 fork)batch import
discussionfeaturehelp wantedreproducibility

Metriche repository

Star
 (26.687 star)
Metriche merge PR
 (Merge medio 9g 15h) (3 PR mergiate in 30 g)

Descrizione

Bug description

If you run the GAN example model https://github.com/Lightning-AI/lightning/blob/master/examples/pl_domain_templates/generative_adversarial_net.py with DDP you are effectively training with only a single GPU because they are all sampling the same latent vector. The offending code is https://github.com/Lightning-AI/lightning/blob/3ff3ec3fdef92fa2f187f06eca41bf08dcc4eb19/examples/pl_domain_templates/generative_adversarial_net.py#L155 to fix this I do something like this in every GAN model I create.

def _augmentation_seed(self):
    if not hasattr(self, 'seeded') or not self.seeded:
        seeds = torch.randint(0, 2**32 - 1, (self.trainer.world_size, ))
        pl.seed_everything(seeds[self.trainer.global_rank], True)
        self.seeded = True

def _sample_latent(self, imgs):
    self._augmentation_seed()
    return torch.randn(imgs.shape[0], self.hparams.latent_dim) 

def training_step(self, batch, batch_idx, optimizer_idx):
        imgs, _ = batch

        # sample noise
        z = self._sample_latent(imgs)
        z = z.type_as(imgs)

It would be nice if lightning could either 1 warn the user about this or 2 (I think this would be better) after all models have been initialized, something like _augmentation_seed is called internally to the Lightning module. I call _augmentation_seed in training_step because I am not 100% sure when all models have been initialized. I also want to point out that this is not only restricted to GANs; it would also affect anyone doing augmentation in the training loop.

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 1.10):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @borda @awaelchli

Guida contributor