parser.add_argument('--noise_aug', type = bool, default = True, help = 'noise_aug') parser.add_argument('--iso', type = int, default = 6400, help = 'noise_level, according to ISO value') parser.add_argument('--random_crop', type = bool, default = True, help = 'random_crop') parser.add_argument('--crop_size', type = int, default = 320, help = 'single patch size') parser.add_argument('--extra_process_train_data', type = bool, default = True, help = 'recover short exposure data') parser.add_argument('--cover_long_exposure', type = bool, default = False, help = 'set long exposure to 0') parser.add_argument('--short_expo_per_pattern', type = int, default = 2, help = 'the number of exposure pixel of 2*2 square') opt = parser.parse_args() # ---------------------------------------- # Test # ---------------------------------------- # Initialize generator = utils.create_generator(opt).cuda() namelist = utils.get_jpgs(opt.in_path_val) test_dataset = dataset.Qbayer2RGB_dataset(opt, 'val', namelist) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = opt.test_batch_size, shuffle = True, num_workers = opt.num_workers, pin_memory = True) sample_folder = os.path.join(opt.val_path, opt.task_name) utils.check_path(sample_folder) # forward val_PSNR = 0 val_SSIM = 0 for i, (in_img, RGBout_img, path) in enumerate(test_loader): # To device # A is for input image, B is for target image in_img = in_img.cuda() RGBout_img = RGBout_img.cuda() #print(path) # Forward propagation
def Trainer_GAN(opt): # ---------------------------------------- # Initialization # ---------------------------------------- # cudnn benchmark cudnn.benchmark = opt.cudnn_benchmark # configurations save_folder = os.path.join(opt.save_path, opt.task_name) sample_folder = os.path.join(opt.sample_path, opt.task_name) if not os.path.exists(save_folder): os.makedirs(save_folder) if not os.path.exists(sample_folder): os.makedirs(sample_folder) # Initialize networks generator = utils.create_generator(opt) discriminator = utils.create_discriminator(opt) perceptualnet = utils.create_perceptualnet() # To device if opt.multi_gpu: generator = nn.DataParallel(generator) generator = generator.cuda() discriminator = nn.DataParallel(discriminator) discriminator = discriminator.cuda() perceptualnet = nn.DataParallel(perceptualnet) perceptualnet = perceptualnet.cuda() else: generator = generator.cuda() discriminator = discriminator.cuda() perceptualnet = perceptualnet.cuda() # ---------------------------------------- # Network dataset # ---------------------------------------- # Handle multiple GPUs gpu_num = torch.cuda.device_count() print("There are %d GPUs used" % gpu_num) opt.train_batch_size *= gpu_num #opt.val_batch_size *= gpu_num opt.num_workers *= gpu_num # Define the dataset train_imglist = utils.get_jpgs(os.path.join(opt.in_path_train)) val_imglist = utils.get_jpgs(os.path.join(opt.in_path_val)) train_dataset = dataset.Qbayer2RGB_dataset(opt, 'train', train_imglist) val_dataset = dataset.Qbayer2RGB_dataset(opt, 'val', val_imglist) print('The overall number of training images:', len(train_imglist)) print('The overall number of validation images:', len(val_imglist)) # Define the dataloader train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.train_batch_size, shuffle=True, num_workers=opt.num_workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=opt.val_batch_size, shuffle=False, num_workers=opt.num_workers, pin_memory=True) # ---------------------------------------- # Network training parameters # ---------------------------------------- # Loss functions criterion_L1 = torch.nn.L1Loss().cuda() class ColorLoss(nn.Module): def __init__(self): super(ColorLoss, self).__init__() self.L1loss = nn.L1Loss() def RGB2YUV(self, RGB): YUV = RGB.clone() YUV[:, 0, :, :] = 0.299 * RGB[:, 0, :, :] + 0.587 * RGB[:, 1, :, :] + 0.114 * RGB[:, 2, :, :] YUV[:, 1, :, :] = -0.14713 * RGB[:, 0, :, :] - 0.28886 * RGB[:, 1, :, :] + 0.436 * RGB[:, 2, :, :] YUV[:, 2, :, :] = 0.615 * RGB[:, 0, :, :] - 0.51499 * RGB[:, 1, :, :] - 0.10001 * RGB[:, 2, :, :] return YUV def forward(self, x, y): yuv_x = self.RGB2YUV(x) yuv_y = self.RGB2YUV(y) return self.L1loss(yuv_x, yuv_y) yuv_loss = ColorLoss() # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr_g, betas=(opt.b1, opt.b2), weight_decay=opt.weight_decay) optimizer_D = torch.optim.Adam(generator.parameters(), lr=opt.lr_d, betas=(opt.b1, opt.b2), weight_decay=opt.weight_decay) # Learning rate decrease def adjust_learning_rate(opt, epoch, iteration, optimizer, lr_gd): # Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs if opt.lr_decrease_mode == 'epoch': lr = lr_gd * (opt.lr_decrease_factor **(epoch // opt.lr_decrease_epoch)) for param_group in optimizer.param_groups: param_group['lr'] = lr if opt.lr_decrease_mode == 'iter': lr = lr_gd * (opt.lr_decrease_factor **(iteration // opt.lr_decrease_iter)) for param_group in optimizer.param_groups: param_group['lr'] = lr # Save the model if pre_train == True def save_model(opt, epoch, iteration, len_dataset, generator): # Define the name of trained model if opt.save_mode == 'epoch': model_name = '%s_gan_noise%.3f_epoch%d_bs%d.pth' % ( opt.net_mode, opt.noise_level, epoch, opt.train_batch_size) if opt.save_mode == 'iter': model_name = '%s_gan_noise%.3f_iter%d_bs%d.pth' % ( opt.net_mode, opt.noise_level, iteration, opt.train_batch_size) save_model_path = os.path.join(opt.save_path, opt.task_name, model_name) # Save model if opt.multi_gpu == True: if opt.save_mode == 'epoch': if (epoch % opt.save_by_epoch == 0) and (iteration % len_dataset == 0): torch.save(generator.module.state_dict(), save_model_path) print( 'The trained model is successfully saved at epoch %d' % (epoch)) if opt.save_mode == 'iter': if iteration % opt.save_by_iter == 0: torch.save(generator.module.state_dict(), save_model_path) print( 'The trained model is successfully saved at iteration %d' % (iteration)) else: if opt.save_mode == 'epoch': if (epoch % opt.save_by_epoch == 0) and (iteration % len_dataset == 0): torch.save(generator.state_dict(), save_model_path) print( 'The trained model is successfully saved at epoch %d' % (epoch)) if opt.save_mode == 'iter': if iteration % opt.save_by_iter == 0: torch.save(generator.state_dict(), save_model_path) print( 'The trained model is successfully saved at iteration %d' % (iteration)) # ---------------------------------------- # Training # ---------------------------------------- # Count start time prev_time = time.time() # Tensorboard writer = SummaryWriter() # For loop training for epoch in range(opt.epochs): # Record learning rate for param_group in optimizer_G.param_groups: writer.add_scalar('data/lr', param_group['lr'], epoch) print('learning rate = ', param_group['lr']) if epoch == 0: iters_done = 0 ### Training for i, (in_img, RGBout_img) in enumerate(train_loader): # To device # A is for input image, B is for target image in_img = in_img.cuda() RGBout_img = RGBout_img.cuda() ## Train Discriminator # Forward propagation out = generator(in_img) optimizer_D.zero_grad() # Fake samples fake_scalar_d = discriminator(in_img, out.detach()) true_scalar_d = discriminator(in_img, RGBout_img) # Overall Loss and optimize loss_D = -torch.mean(true_scalar_d) + torch.mean(fake_scalar_d) loss_D.backward() #torch.nn.utils.clip_grad_norm(discriminator.parameters(), opt.grad_clip_norm) optimizer_D.step() ## Train Generator # Forward propagation out = generator(in_img) # GAN loss fake_scalar = discriminator(in_img, out) L_gan = -torch.mean(fake_scalar) * opt.lambda_gan # Perceptual loss features fake_B_fea = perceptualnet(utils.normalize_ImageNet_stats(out)) true_B_fea = perceptualnet( utils.normalize_ImageNet_stats(RGBout_img)) L_percep = opt.lambda_percep * criterion_L1(fake_B_fea, true_B_fea) # Pixel loss L_pixel = opt.lambda_pixel * criterion_L1(out, RGBout_img) # Color loss L_color = opt.lambda_color * yuv_loss(out, RGBout_img) # Sum up to total loss loss = L_pixel + L_percep + L_gan + L_color # Record losses writer.add_scalar('data/L_pixel', L_pixel.item(), iters_done) writer.add_scalar('data/L_percep', L_percep.item(), iters_done) writer.add_scalar('data/L_color', L_color.item(), iters_done) writer.add_scalar('data/L_gan', L_gan.item(), iters_done) writer.add_scalar('data/L_total', loss.item(), iters_done) writer.add_scalar('data/loss_D', loss_D.item(), iters_done) # Backpropagate gradients optimizer_G.zero_grad() loss.backward() #torch.nn.utils.clip_grad_norm(generator.parameters(), opt.grad_clip_norm) optimizer_G.step() # Determine approximate time left iters_done = epoch * len(train_loader) + i + 1 iters_left = opt.epochs * len(train_loader) - iters_done time_left = datetime.timedelta(seconds=iters_left * (time.time() - prev_time)) prev_time = time.time() # Print log print( "\r[Epoch %d/%d] [Batch %d/%d] [Total Loss: %.4f] [L_pixel: %.4f]" % ((epoch + 1), opt.epochs, i, len(train_loader), loss.item(), L_pixel.item())) print( "\r[L_percep: %.4f] [L_color: %.4f] [L_gan: %.4f] [loss_D: %.4f] Time_left: %s" % (L_percep.item(), L_color.item(), L_gan.item(), loss_D.item(), time_left)) # Save model at certain epochs or iterations save_model(opt, (epoch + 1), iters_done, len(train_loader), generator) # Learning rate decrease at certain epochs adjust_learning_rate(opt, (epoch + 1), iters_done, optimizer_G, opt.lr_g) adjust_learning_rate(opt, (epoch + 1), iters_done, optimizer_D, opt.lr_d) ### Sample data every epoch if (epoch + 1) % 1 == 0: img_list = [out, RGBout_img] name_list = ['pred', 'gt'] utils.save_sample_png(sample_folder=sample_folder, sample_name='train_epoch%d' % (epoch + 1), img_list=img_list, name_list=name_list, pixel_max_cnt=255) ### Validation val_PSNR = 0 num_of_val_image = 0 for j, (in_img, RGBout_img) in enumerate(val_loader): # To device # A is for input image, B is for target image in_img = in_img.cuda() RGBout_img = RGBout_img.cuda() # Forward propagation with torch.no_grad(): out = generator(in_img) # Accumulate num of image and val_PSNR num_of_val_image += in_img.shape[0] val_PSNR += utils.psnr(out, RGBout_img, 1) * in_img.shape[0] val_PSNR = val_PSNR / num_of_val_image ### Sample data every epoch if (epoch + 1) % 1 == 0: img_list = [out, RGBout_img] name_list = ['pred', 'gt'] utils.save_sample_png(sample_folder=sample_folder, sample_name='val_epoch%d' % (epoch + 1), img_list=img_list, name_list=name_list, pixel_max_cnt=255) # Record average PSNR writer.add_scalar('data/val_PSNR', val_PSNR, epoch) print('PSNR at epoch %d: %.4f' % ((epoch + 1), val_PSNR)) writer.close()