Lightning-AI/pytorch-lightning

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

Open

#16 166 ouverte le 21 déc. 2022

Voir sur GitHub
 (3 commentaires) (2 réactions) (0 assignés)Python (3 233 forks)batch import
discussionfeaturehelp wantedreproducibility

Métriques du dépôt

Stars
 (26 687 stars)
Métriques de merge PR
 (Merge moyen 9j 15h) (3 PRs mergées en 30 j)

Description

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

Guide contributeur