torch.randn() + DDP + GANs are easy to get wrong with lightning
#16,166 opened on 2022年12月21日
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