(critic, real, fake, device)
| 13 | |
| 14 | |
| 15 | def gradient_penalty(critic, real, fake, device): |
| 16 | BATCH_SIZE, C, H, W = real.shape |
| 17 | alpha = torch.rand((BATCH_SIZE, 1, 1, 1)).repeat(1, C, H, W).to(device) |
| 18 | interpolated_images = real * alpha + fake.detach() * (1 - alpha) |
| 19 | interpolated_images.requires_grad_(True) |
| 20 | |
| 21 | # Calculate critic scores |
| 22 | mixed_scores = critic(interpolated_images) |
| 23 | |
| 24 | # Take the gradient of the scores with respect to the images |
| 25 | gradient = torch.autograd.grad( |
| 26 | inputs=interpolated_images, |
| 27 | outputs=mixed_scores, |
| 28 | grad_outputs=torch.ones_like(mixed_scores), |
| 29 | create_graph=True, |
| 30 | retain_graph=True, |
| 31 | )[0] |
| 32 | gradient = gradient.view(gradient.shape[0], -1) |
| 33 | gradient_norm = gradient.norm(2, dim=1) |
| 34 | gradient_penalty = torch.mean((gradient_norm - 1) ** 2) |
| 35 | return gradient_penalty |
| 36 | |
| 37 | |
| 38 | def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"): |
nothing calls this directly
no outgoing calls
no test coverage detected