| 18 | torch.backends.cudnn.benchmark = True |
| 19 | |
| 20 | def train_fn(loader, disc, gen, opt_gen, opt_disc, mse, bce, vgg_loss,epoch): |
| 21 | |
| 22 | loop = tqdm(loader, leave=True) |
| 23 | for idx, (low_res, high_res) in enumerate(loop): |
| 24 | high_res = high_res.to(config.DEVICE) |
| 25 | low_res = low_res.to(config.DEVICE) |
| 26 | |
| 27 | ### Train Discriminator: max log(D(x)) + log(1 - D(G(z))) |
| 28 | # print('low_res.shape: {}'.format(low_res.shape)) |
| 29 | # print('high_res.shape: {}'.format(high_res.shape)) |
| 30 | fake = gen(low_res) |
| 31 | disc_real = disc(high_res) |
| 32 | disc_fake = disc(fake.detach()) |
| 33 | # print('disc_fake.shape: {}'.format(disc_fake.shape)) |
| 34 | # print('disc_real.shape: {}'.format(disc_real.shape)) |
| 35 | disc_loss_real = bce( |
| 36 | disc_real, torch.ones_like(disc_real) - 0.1 * torch.rand_like(disc_real) |
| 37 | ) |
| 38 | disc_loss_fake = bce(disc_fake, torch.zeros_like(disc_fake)) |
| 39 | loss_disc = disc_loss_fake + disc_loss_real |
| 40 | |
| 41 | opt_disc.zero_grad() |
| 42 | loss_disc.backward() |
| 43 | opt_disc.step() |
| 44 | |
| 45 | # Train Generator: min log(1 - D(G(z))) <-> max log(D(G(z)) |
| 46 | disc_fake = disc(fake) |
| 47 | # l2_loss = mse(fake, high_res) |
| 48 | adversarial_loss = 1e-3 * bce(disc_fake, torch.ones_like(disc_fake)) |
| 49 | #0.006 = 1 / (w x h) = 1 / (14 x 14) |
| 50 | loss_for_vgg = 0.006 * vgg_loss(fake, high_res) |
| 51 | gen_loss = loss_for_vgg + adversarial_loss |
| 52 | |
| 53 | opt_gen.zero_grad() |
| 54 | gen_loss.backward() |
| 55 | opt_gen.step() |
| 56 | |
| 57 | if idx % 200 == 0: |
| 58 | plot_examples("test_images/", gen) |
| 59 | |
| 60 | loop.set_postfix( |
| 61 | epoch = epoch, |
| 62 | loss_critic=loss_disc.item(), |
| 63 | loss_gen=gen_loss.item() |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def main(): |