Lightning-AI/pytorch-lightning

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

Open

#16,166 opened on 2022年12月21日

GitHub で見る
 (3 comments) (2 reactions) (0 assignees)Python (3,233 forks)batch import
discussionfeaturehelp wantedreproducibility

Repository metrics

Stars
 (26,687 stars)
PR merge metrics
 (平均マージ 9d 15h) (30d で 3 merged PRs)

説明

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

コントリビューターガイド