| 16 | |
| 17 | # Print losses occasionally and print to tensorboard |
| 18 | def plot_to_tensorboard( |
| 19 | writer, loss_critic, loss_gen, real, fake, tensorboard_step |
| 20 | ): |
| 21 | writer.add_scalar("Loss Critic", loss_critic, global_step=tensorboard_step) |
| 22 | |
| 23 | with torch.no_grad(): |
| 24 | # take out (up to) 8 examples to plot |
| 25 | img_grid_real = torchvision.utils.make_grid(real[:8], normalize=True) |
| 26 | img_grid_fake = torchvision.utils.make_grid(fake[:8], normalize=True) |
| 27 | writer.add_image("Real", img_grid_real, global_step=tensorboard_step) |
| 28 | writer.add_image("Fake", img_grid_fake, global_step=tensorboard_step) |
| 29 | save_image(img_grid_fake,os.path.join('images','fake_'+str(tensorboard_step) + '.png')) |
| 30 | save_image(img_grid_real,os.path.join('images','real_'+str(tensorboard_step) + '.png')) |
| 31 | |
| 32 | |
| 33 | def gradient_penalty(critic, real, fake, alpha, train_step, device="cpu"): |