Lightning-AI/pytorch-lightning

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

Open

#16.166 aberto em 21 de dez. de 2022

Ver no GitHub
 (3 comments) (2 reactions) (0 assignees)Python (3.233 forks)batch import
discussionfeaturehelp wantedreproducibility

Métricas do repositório

Stars
 (26.687 stars)
Métricas de merge de PR
 (Mesclagem média 9d 15h) (3 fundiu PRs em 30d)

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

Guia do colaborador