def train(args): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") transform = transforms.Compose( [crop(args.scale, args.patch_size), augmentation()]) dataset = mydata(GT_path=args.GT_path, LR_path=args.LR_path, in_memory=args.in_memory, transform=transform) loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) generator = Generator(img_feat=3, n_feats=64, kernel_size=3, num_block=args.res_num, scale=args.scale) if args.fine_tuning: generator.load_state_dict(torch.load(args.generator_path)) print("pre-trained model is loaded") print("path : %s" % (args.generator_path)) generator = generator.to(device) generator.train() l2_loss = nn.MSELoss() g_optim = optim.Adam(generator.parameters(), lr=1e-4) pre_epoch = 0 fine_epoch = 0 #### Train using L2_loss while pre_epoch < args.pre_train_epoch: for i, tr_data in enumerate(loader): gt = tr_data['GT'].to(device) lr = tr_data['LR'].to(device) output, _ = generator(lr) loss = l2_loss(gt, output) g_optim.zero_grad() loss.backward() g_optim.step() pre_epoch += 1 if pre_epoch % 2 == 0: print(pre_epoch) print(loss.item()) print('=========') if pre_epoch % 800 == 0: torch.save(generator.state_dict(), './model/pre_trained_model_%03d.pt' % pre_epoch) #### Train using perceptual & adversarial loss vgg_net = vgg19().to(device) vgg_net = vgg_net.eval() discriminator = Discriminator(patch_size=args.patch_size * args.scale) discriminator = discriminator.to(device) discriminator.train() d_optim = optim.Adam(discriminator.parameters(), lr=1e-4) scheduler = optim.lr_scheduler.StepLR(g_optim, step_size=2000, gamma=0.1) VGG_loss = perceptual_loss(vgg_net) cross_ent = nn.BCELoss() tv_loss = TVLoss() real_label = torch.ones((args.batch_size, 1)).to(device) fake_label = torch.zeros((args.batch_size, 1)).to(device) while fine_epoch < args.fine_train_epoch: scheduler.step() for i, tr_data in enumerate(loader): gt = tr_data['GT'].to(device) lr = tr_data['LR'].to(device) ## Training Discriminator output, _ = generator(lr) fake_prob = discriminator(output) real_prob = discriminator(gt) d_loss_real = cross_ent(real_prob, real_label) d_loss_fake = cross_ent(fake_prob, fake_label) d_loss = d_loss_real + d_loss_fake g_optim.zero_grad() d_optim.zero_grad() d_loss.backward() d_optim.step() ## Training Generator output, _ = generator(lr) fake_prob = discriminator(output) _percep_loss, hr_feat, sr_feat = VGG_loss((gt + 1.0) / 2.0, (output + 1.0) / 2.0, layer=args.feat_layer) L2_loss = l2_loss(output, gt) percep_loss = args.vgg_rescale_coeff * _percep_loss adversarial_loss = args.adv_coeff * cross_ent( fake_prob, real_label) total_variance_loss = args.tv_loss_coeff * tv_loss( args.vgg_rescale_coeff * (hr_feat - sr_feat)**2) g_loss = percep_loss + adversarial_loss + total_variance_loss + L2_loss g_optim.zero_grad() d_optim.zero_grad() g_loss.backward() g_optim.step() fine_epoch += 1 if fine_epoch % 2 == 0: print(fine_epoch) print(g_loss.item()) print(d_loss.item()) print('=========') if fine_epoch % 500 == 0: torch.save(generator.state_dict(), './model/SRGAN_gene_%03d.pt' % fine_epoch) torch.save(discriminator.state_dict(), './model/SRGAN_discrim_%03d.pt' % fine_epoch)
'ssims'] / valing_results['batch_sizes'] sum_val_psnr += valing_results['psnr'] sum_val_ssim += valing_results['ssim'] # save model parameters: if not os.path.exists(epoch_path): os.makedirs(epoch_path) # save model every 25 / 50 epochs if epoch % v == 0 and epoch != 0: torch.save( G.state_dict(), epoch_path + "/" + 'netG_epoch_%d_%d.pth' % (upscale_factor, epoch)) torch.save( D.state_dict(), epoch_path + "/" + 'netD_epoch_%d_%d.pth' % (upscale_factor, epoch)) # save loss, scores, psnr, ssim results['d_loss'].append(running_results['d_loss'] / running_results['batch_sizes']) results['g_loss'].append(running_results['g_loss'] / running_results['batch_sizes']) results['d_score'].append(running_results['d_score'] / running_results['batch_sizes']) results['g_score'].append(running_results['g_score'] / running_results['batch_sizes']) results['psnr'].append(valing_results['psnr']) results['ssim'].append(valing_results['ssim']) avg_psnr = sum_val_psnr / batch_Size