| 173 | self.pl_no_weight_grad = pl_no_weight_grad |
| 174 | |
| 175 | def cal_pl_reg(self, fake_images, ws): |
| 176 | #ws refers to weight style |
| 177 | #receives new fake_images of original batch (in original implementation, fakes_images used for calculating g_loss and pl_loss is generated independently) |
| 178 | pl_noise = torch.randn_like(fake_images) / np.sqrt(fake_images.shape[2] * fake_images.shape[3]) |
| 179 | with conv2d_gradfix.no_weight_gradients(self.pl_no_weight_grad): |
| 180 | pl_grads = torch.autograd.grad(outputs=[(fake_images * pl_noise).sum()], inputs=[ws], create_graph=True, only_inputs=True)[0] |
| 181 | pl_lengths = pl_grads.square().sum(2).mean(1).sqrt() |
| 182 | pl_mean = self.pl_mean.lerp(pl_lengths.mean(), self.pl_decay) |
| 183 | self.pl_mean.copy_(pl_mean.detach()) |
| 184 | pl_penalty = (pl_lengths - pl_mean).square() |
| 185 | loss_Gpl = (pl_penalty * self.pl_weight).mean(0) |
| 186 | return loss_Gpl |
| 187 | |
| 188 | |
| 189 | def enable_allreduce(dict_): |