| 212 | |
| 213 | # Print losses occasionally and print to tensorboard |
| 214 | def plot_to_tensorboard( |
| 215 | writer, loss_critic, loss_gen, real, fake, tensorboard_step |
| 216 | ): |
| 217 | writer.add_scalar("Loss Critic", loss_critic, global_step=tensorboard_step) |
| 218 | |
| 219 | with torch.no_grad(): |
| 220 | # take out (up to) 8 examples to plot |
| 221 | img_grid_real = torchvision.utils.make_grid(real[:8], normalize=True) |
| 222 | img_grid_fake = torchvision.utils.make_grid(fake[:8], normalize=True) |
| 223 | writer.add_image("Real", img_grid_real, global_step=tensorboard_step) |
| 224 | writer.add_image("Fake", img_grid_fake, global_step=tensorboard_step) |
| 225 | save_image(img_grid_fake,os.path.join('images','fake_'+tensorboard_step + '.png')) |
| 226 | save_image(img_grid_real,os.path.join('images','real_'+tensorboard_step + '.png')) |
| 227 | |
| 228 | |
| 229 | def gradient_penalty(critic, real, fake, alpha, train_step, device="cpu"): |