| # PyTorch StudioGAN: https://github.com/POSTECH-CVLab/PyTorch-StudioGAN | |
| # The MIT License (MIT) | |
| # See license file or visit https://github.com/POSTECH-CVLab/PyTorch-StudioGAN for details | |
| # src/utils/apa_aug.py | |
| import torch | |
| def apply_apa_aug(real_images, fake_images, apa_p, local_rank): | |
| # Apply Adaptive Pseudo Augmentation (APA) | |
| # https://github.com/EndlessSora/DeceiveD/blob/main/training/loss.py | |
| batch_size = real_images.shape[0] | |
| pseudo_flag = torch.ones([batch_size, 1, 1, 1], device=local_rank) | |
| pseudo_flag = torch.where(torch.rand([batch_size, 1, 1, 1], device=local_rank) < apa_p, | |
| pseudo_flag, torch.zeros_like(pseudo_flag)) | |
| if torch.allclose(pseudo_flag, torch.zeros_like(pseudo_flag)): | |
| return real_images | |
| else: | |
| assert fake_images is not None | |
| return fake_images * pseudo_flag + real_images * (1 - pseudo_flag) | |