def main(): # Split the dataset train_dataset = sunnerData.ImageDataset( root=[['/home/sunner/Music/waiting_for_you_dataset/wait'], ['/home/sunner/Music/waiting_for_you_dataset/real_world']], transform=None, split_ratio=0.1, save_file=True) del train_dataset test_dataset = sunnerData.ImageDataset( file_name='.split.pkl', transform=transforms.Compose([ sunnertransforms.Resize((160, 320)), sunnertransforms.ToTensor(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ])) # Create the data loader loader = sunnerData.DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=2) # Use upper wrapper to assign particular iteration loader = sunnerData.IterationLoader(loader, max_iter=1) # Show! for batch_img, _ in loader: batch_img = sunnertransforms.asImg(batch_img, size=(160, 320)) cv2.imshow('show_window', batch_img[0][:, :, ::-1]) cv2.waitKey(0)
def main(): # Define op first transform_op = transforms.Compose([ sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ]) # Define loader loader = sunnerData.DataLoader( sunnerData.VideoDataset( root = [ ['/home/sunner/Dataset/flower/A'], ['/home/sunner/Dataset/flower/B'] ], transform = transform_op, T = 20 ), batch_size=2, shuffle=False, num_workers = 2 ) # Looping for _, seq in loader: seq = [_.squeeze(1) for _ in torch.chunk(seq, seq.size(1), dim = 1)] # BTCHW -> T * BCHW for i in range(10): # repeat for 10 times for img in seq: batch_img = sunnertransforms.asImg(img, size = (320, 640)) cv2.imshow('show_window', batch_img[0][:, :, ::-1]) cv2.waitKey(10) break
def main(): # Define the loader to generate the pallete object loader = sunnerData.DataLoader( sunnerData.ImageDataset( root=[tag_folder], transform=transforms.Compose([ sunnertransforms.ToTensor(), ]), save_file=False # Don't save the record file, be careful! ), batch_size=2, shuffle=False, num_workers=2) pallete = sunnertransforms.getCategoricalMapping(loader, path='pallete.json')[0] del loader # Define the actual loader loader = sunnerData.DataLoader(sunnerData.ImageDataset( root=[img_folder, tag_folder], transform=transforms.Compose([ sunnertransforms.Resize((512, 1024)), sunnertransforms.ToTensor(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ])), batch_size=32, shuffle=False, num_workers=2) # Define the reverse operator goto_op = sunnertransforms.CategoricalTranspose( pallete=pallete, direction=sunnertransforms.COLOR2ONEHOT) back_op = sunnertransforms.CategoricalTranspose( pallete=pallete, direction=sunnertransforms.ONEHOT2COLOR) # Show! for _, batch_index in loader: batch_img = back_op(goto_op(batch_index)) batch_img = sunnertransforms.asImg(batch_img, size=(512, 1024)) cv2.imshow('show_window', batch_img[0][:, :, ::-1]) cv2.waitKey(0) break
def main(): # Define the loader to generate the pallete object loader = sunnerData.DataLoader(sunnerData.ImageDataset( root = [ tag_folder ], transforms = transforms.Compose([ sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.UnNormalize(mean=[0, 0, 0], std=[255, 255, 255]), # Remember to transfer back to [0~255] before generate pallete sunnertransforms.Transpose(sunnertransforms.BCHW2BHWC) # Remember to transfer back to BHWC before generate pallete ]) ), batch_size = 2, shuffle = False, num_workers = 2 ) pallete = sunnertransforms.getCategoricalMapping(loader, path = 'pallete.json')[0] del loader # Define the actual loader loader = sunnerData.DataLoader(sunnerData.ImageDataset( root = [ img_folder, tag_folder ], transforms = transforms.Compose([ sunnertransforms.Resize((512, 1024)), sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Normalize(mean = [0.5, 0.5, 0.5], std = [0.5, 0.5, 0.5]), ])), batch_size = 32, shuffle = False, num_workers = 2 ) # Define the reverse operator goto_op = sunnertransforms.CategoricalTranspose(pallete = pallete, direction = sunnertransforms.COLOR2ONEHOT) back_op = sunnertransforms.CategoricalTranspose(pallete = pallete, direction = sunnertransforms.ONEHOT2COLOR) # Show! for _, batch_index in loader: batch_img = back_op(goto_op(batch_index)) batch_img = sunnertransforms.asImg(batch_img, size = (512, 1024)) cv2.imshow('show_window', batch_img[0][:, :, ::-1]) cv2.waitKey(0) break
def main(): # Create the fundemental data loader loader = sunnerData.DataLoader(sunnerData.ImageDataset( root=[['/home/sunner/Music/waiting_for_you_dataset/wait'], ['/home/sunner/Music/waiting_for_you_dataset/real_world']], transform=transforms.Compose([ sunnertransforms.Resize((160, 320)), sunnertransforms.ToTensor(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ])), batch_size=32, shuffle=False, num_workers=2) # Use upper wrapper to assign particular iteration loader = sunnerData.IterationLoader(loader, max_iter=1) # Show! for batch_img, _ in loader: batch_img = sunnertransforms.asImg(batch_img, size=(160, 320)) cv2.imshow('show_window', batch_img[0][:, :, ::-1]) cv2.waitKey(0)
def inference(opts): # Load the image ops = transforms.Compose([ sunnertransforms.Resize((opts.H, opts.W)), sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ]) img1 = cv2.imread(opts.image1) img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2YCrCb) img1 = torch.unsqueeze(ops(img1), 0) img2 = cv2.imread(opts.image2) img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2YCrCb) img2 = torch.unsqueeze(ops(img2), 0) # Load the pre-trained model model = DeepFuse() state = torch.load(opts.model) model.load_state_dict(state['model']) model.to(opts.device) model.eval() criterion = MEF_SSIM_Loss().to(opts.device) # Fuse! with torch.no_grad(): # Forward img1, img2 = img1.to(opts.device), img2.to(opts.device) img1_lum = img1[:, 0:1] img2_lum = img2[:, 0:1] model.setInput(img1_lum, img2_lum) y_f = model.forward() _, y_hat = criterion(y_1=img1_lum, y_2=img2_lum, y_f=y_f) # Save the image img = fusePostProcess(y_f, y_hat, img1, img2, single=False) cv2.imwrite(opts.res, img[0, :, :, :])
def evalModel(args, model): # Create data loader loader = sunnerData.ImageLoader(sunnerData.ImageDataset( root_list=[args.folder_path, args.mask_path], transform=transforms.Compose([ sunnertransforms.Rescale((args.size, args.size)), sunnertransforms.ToTensor(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize() ]), sample_method=sunnerData.OVER_SAMPLING), batch_size=1, shuffle=False, num_workers=2) # Compute the PSNR and record psnr_list = [] bar = tqdm(loader) for image, mask in bar: # Double the tensor to adapt with BN image = torch.cat([image, image], 0) mask = torch.cat([mask, mask], 0) # forward mask = (mask + 1) / 2 model.setInput(target=image, mask=mask) model.forward() _, recon_img, _ = model.getOutput() psnr = compare_psnr(image[0].detach().cpu().numpy(), recon_img[0].detach().cpu().numpy()) psnr_list.append(psnr) # Show the result print('\n\n') print('-' * 20, 'Complete evaluation', '-' * 20) print('Testing average psnr: %.4f' % np.mean(psnr_list))
def main(opts): # Create the data loader loader = sunnerData.DataLoader( sunnerData.ImageDataset(root=[[opts.path]], transform=transforms.Compose([ sunnertransforms.Resize((1024, 1024)), sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose( sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ])), batch_size=opts.batch_size, shuffle=True, ) # Create the model start_epoch = 0 G = StyleGenerator() D = StyleDiscriminator() # Load the pre-trained weight if os.path.exists(opts.resume): INFO("Load the pre-trained weight!") state = torch.load(opts.resume) G.load_state_dict(state['G']) D.load_state_dict(state['D']) start_epoch = state['start_epoch'] else: INFO( "Pre-trained weight cannot load successfully, train from scratch!") # Multi-GPU support if torch.cuda.device_count() > 1: INFO("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs") G = nn.DataParallel(G) D = nn.DataParallel(D) G.to(opts.device) D.to(opts.device) # Create the criterion, optimizer and scheduler optim_D = optim.Adam(D.parameters(), lr=0.00001, betas=(0.5, 0.999)) optim_G = optim.Adam(G.parameters(), lr=0.00001, betas=(0.5, 0.999)) scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99) scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99) # Train fix_z = torch.randn([opts.batch_size, 512]).to(opts.device) softplus = nn.Softplus() Loss_D_list = [0.0] Loss_G_list = [0.0] for ep in range(start_epoch, opts.epoch): bar = tqdm(loader) loss_D_list = [] loss_G_list = [] for i, (real_img, ) in enumerate(bar): # ======================================================================================================= # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) # ======================================================================================================= # Compute adversarial loss toward discriminator D.zero_grad() real_img = real_img.to(opts.device) real_logit = D(real_img) fake_img = G(torch.randn([real_img.size(0), 512]).to(opts.device)) fake_logit = D(fake_img.detach()) d_loss = softplus(fake_logit).mean() d_loss = d_loss + softplus(-real_logit).mean() if opts.r1_gamma != 0.0: r1_penalty = R1Penalty(real_img.detach(), D) d_loss = d_loss + r1_penalty * (opts.r1_gamma * 0.5) if opts.r2_gamma != 0.0: r2_penalty = R2Penalty(fake_img.detach(), D) d_loss = d_loss + r2_penalty * (opts.r2_gamma * 0.5) loss_D_list.append(d_loss.item()) # Update discriminator d_loss.backward() optim_D.step() # ======================================================================================================= # (2) Update G network: maximize log(D(G(z))) # ======================================================================================================= if i % CRITIC_ITER == 0: G.zero_grad() fake_logit = D(fake_img) g_loss = softplus(-fake_logit).mean() loss_G_list.append(g_loss.item()) # Update generator g_loss.backward() optim_G.step() # Output training stats bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format( ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1])) # Save the result Loss_G_list.append(np.mean(loss_G_list)) Loss_D_list.append(np.mean(loss_D_list)) # Check how the generator is doing by saving G's output on fixed_noise with torch.no_grad(): fake_img = G(fix_z).detach().cpu() save_image(fake_img, os.path.join(opts.det, 'images', str(ep) + '.png'), nrow=4, normalize=True) # Save model state = { 'G': G.state_dict(), 'D': D.state_dict(), 'Loss_G': Loss_G_list, 'Loss_D': Loss_D_list, 'start_epoch': ep, } torch.save(state, os.path.join(opts.det, 'models', 'latest.pth')) scheduler_D.step() scheduler_G.step() # Plot the total loss curve Loss_D_list = Loss_D_list[1:] Loss_G_list = Loss_G_list[1:] plotLossCurve(opts, Loss_D_list, Loss_G_list)
def train(opts): def log(string, name="stylegan.log"): with open(name, 'a') as f: f.write(string + '\n') writer = SummaryWriter(str(opts.output)) loader = dataset.DataLoader(dataset=dataset.ImageDataset( [[opts.input]], transform=transforms.Compose([ trans.Resize((opts.imsize, opts.imsize)), trans.ToTensor(), trans.ToFloat(), trans.Transpose(trans.BHWC2BCHW), trans.Normalize() ])), batch_size=opts.batch_size, shuffle=True) G = opts.G D = opts.D step = 0 start_epoch = opts.start_epoch if opts.resume: try: assert os.path.exists(opts.resume) state = torch.load(opts.resume) G.load_state_dict(state['G']) D.load_state_dict(state['D']) start_epoch = state['start_epoch'] logger.info("Load Pretrained Weight") except: logger.warn("Resume Files cannot Load") logger.info("Train from Scratch") else: logger.info("Train from Scratch") if torch.cuda.device_count() > 1 and opts.device == 'cuda': logger.info(f"{torch.cuda.device_count()} GPUs found.") G = nn.DataParallel(G) D = nn.DataParallel(D) if opts.device == 'cuda': torch.backends.cudnn.benchmark = True G.to(opts.device) D.to(opts.device) optimG = Adam(G.parameters(), lr=opts.g_lr, betas=opts.betas) optimD = Adam(D.parameters(), lr=opts.d_lr, betas=opts.betas) schedulerG = lr_scheduler.ExponentialLR(optimG, gamma=opts.g_lrdecay) schedulerD = lr_scheduler.ExponentialLR(optimD, gamma=opts.d_lrdecay) fixed_z = torch.randn([opts.batch_size, 512]).to(opts.device) sp = nn.Softplus() g_store = [0.0] d_store = [0.0] for epoch in range(start_epoch, opts.epochs + 1): bar = tqdm(loader) glosses = [] dlosses = [] for i, (real, ) in enumerate(bar): step += 1 D.zero_grad() real = real.to(opts.device) Dr = D(real) writer.add_graph(D, real) z = torch.randn([real.size(0), 512]).to(opts.device) fake = G(z) writer.add_graph(G, z) Df = D(fake.detach()) Dloss = sp(Df).mean() + sp(-Dr).mean() if opts.r1gamma > 0: r1 = r1_penalty(real.detach(), D) Dloss = Dloss + r1 * (opts.r1gamma * .5) if opts.r2gamma > 0: r2 = r2_penalty(fake.detach(), D) Dloss = Dloss + r2 * (opts.r2gamma * .5) dlosses.append(Dloss.item()) Dloss.backward() optimD.step() if i % opts.critic_iters == 0: G.zero_grad() Df = D(fake) Gloss = sp(-Df).mean() glosses.append(Gloss.item()) Gloss.backward() optimG.step() if i % opts.show_interval == 0: with torch.no_grad(): nr = int(math.ceil(math.sqrt(opts.batch_size))) z = torch.randn([real.size(0), 512]).to(opts.device) img = G(z) save_image(img.detach().cpu(), os.path.join(opts.output, 'images', 'normal', f'{epoch:04}_{i:06}.png'), nrow=nr, normalize=True) fakes = utils.make_grid(img, nr, padding=0) fakes = fakes.to(torch.float32).cpu().numpy() fakes = np.clip((fakes / 2) + 0.5, 0, 1) writer.add_image(f"EPOCH{epoch}/Random", torch.from_numpy(fakes), i) img = G(fixed_z) save_image(img.detach().cpu(), os.path.join(opts.output, 'images', 'fixed', f'{epoch:04}_{i:06}.png'), nrow=nr, normalize=True) fakes = utils.make_grid(img, nr, padding=0) fakes = fakes.to(torch.float32).cpu().numpy() fakes = np.clip((fakes / 2) + 0.5, 0, 1) writer.add_image(f"EPOCH{epoch}/Fixed", torch.from_numpy(fakes), i) writer.add_scalar(f"LOSS/Generator", Gloss.item(), global_step=step) writer.add_scalar(f"LOSS/Discriminator", Dloss.item(), global_step=step) bar.set_description( f"Epoch {epoch}/{opts.epochs} G: {glosses[-1]:.6f} D: {dlosses[-1]:.6f}" ) g_store.append(np.mean(glosses)) d_store.append(np.mean(dlosses)) state = { 'G': G.state_dict(), 'D': D.state_dict(), 'Loss_G': g_store, 'Loss_D': d_store, 'start_epoch': epoch, 'opts': opts } torch.save(state, os.path.join(opts.output, 'models', 'latest.pth')) if epoch % 10 == 0: torch.save(state, os.path.join(opts.output, 'models', f'{epoch:04}.pth')) schedulerD.step() schedulerG.step()
def main(opts): # Create the data loader loader = sunnerData.DataLoader(sunnerData.ImageDataset( root=[[opts.path]], transform=transforms.Compose([ sunnertransforms.Resize((128, 128)), sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ])), batch_size=opts.batch_size, shuffle=True, num_workers=4) # Create the model if opts.type == 'style': G = StyleGenerator().to(opts.device) else: G = Generator().to(opts.device) D = Discriminator().to(opts.device) # Load the pre-trained weight if os.path.exists(opts.resume): INFO("Load the pre-trained weight!") state = torch.load(opts.resume) G.load_state_dict(state['G']) D.load_state_dict(state['D']) else: INFO( "Pre-trained weight cannot load successfully, train from scratch!") # Create the criterion, optimizer and scheduler optim_D = optim.Adam(D.parameters(), lr=0.0001, betas=(0.5, 0.999)) optim_G = optim.Adam(G.parameters(), lr=0.0001, betas=(0.5, 0.999)) scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99) scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99) # Train fix_z = torch.randn([opts.batch_size, 512]).to(opts.device) Loss_D_list = [0.0] Loss_G_list = [0.0] for ep in range(opts.epoch): bar = tqdm(loader) loss_D_list = [] loss_G_list = [] for i, (real_img, ) in enumerate(bar): # ======================================================================================================= # Update discriminator # ======================================================================================================= # Compute adversarial loss toward discriminator real_img = real_img.to(opts.device) real_logit = D(real_img) fake_img = G(torch.randn([real_img.size(0), 512]).to(opts.device)) fake_logit = D(fake_img.detach()) d_loss = -(real_logit.mean() - fake_logit.mean()) + gradient_penalty( real_img.data, fake_img.data, D) * 10.0 loss_D_list.append(d_loss.item()) # Update discriminator optim_D.zero_grad() d_loss.backward() optim_D.step() # ======================================================================================================= # Update generator # ======================================================================================================= if i % CRITIC_ITER == 0: # Compute adversarial loss toward generator fake_img = G( torch.randn([opts.batch_size, 512]).to(opts.device)) fake_logit = D(fake_img) g_loss = -fake_logit.mean() loss_G_list.append(g_loss.item()) # Update generator D.zero_grad() optim_G.zero_grad() g_loss.backward() optim_G.step() bar.set_description(" {} [G]: {} [D]: {}".format( ep, loss_G_list[-1], loss_D_list[-1])) # Save the result Loss_G_list.append(np.mean(loss_G_list)) Loss_D_list.append(np.mean(loss_D_list)) fake_img = G(fix_z) save_image(fake_img, os.path.join(opts.det, 'images', str(ep) + '.png'), nrow=4, normalize=True) state = { 'G': G.state_dict(), 'D': D.state_dict(), 'Loss_G': Loss_G_list, 'Loss_D': Loss_D_list, } torch.save(state, os.path.join(opts.det, 'models', 'latest.pth')) scheduler_D.step() scheduler_G.step() # Plot the total loss curve Loss_D_list = Loss_D_list[1:] Loss_G_list = Loss_G_list[1:] plotLossCurve(opts, Loss_D_list, Loss_G_list)
def main(opts): # Create the data loader loader = sunnerData.DataLoader(sunnerData.ImageDataset( root=[[opts.path]], transform=transforms.Compose([ sunnertransforms.Resize((opts.resolution, opts.resolution)), sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ])), batch_size=opts.batch_size, shuffle=True, drop_last=True ) # Create the model start_epoch = 0 G = G_stylegan2(fmap_base=opts.fmap_base, resolution=opts.resolution, mapping_layers=opts.mapping_layers, opts=opts, return_dlatents=True) D = D_stylegan2(fmap_base=opts.fmap_base, resolution=opts.resolution, structure='resnet') # Load the pre-trained weight if os.path.exists(opts.resume): INFO("Load the pre-trained weight!") state = torch.load(opts.resume) G.load_state_dict(state['G']) D.load_state_dict(state['D']) start_epoch = state['start_epoch'] else: INFO("Pre-trained weight cannot load successfully, train from scratch!") # Multi-GPU support if torch.cuda.device_count() > 1: INFO("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs") G = torch.nn.DataParallel(G) D = torch.nn.DataParallel(D) G.to(opts.device) D.to(opts.device) # Create the criterion, optimizer and scheduler lr_D = 0.0015 lr_G = 0.0015 optim_D = torch.optim.Adam(D.parameters(), lr=lr_D, betas=(0.9, 0.999)) # g_mapping has 100x lower learning rate params_G = [{"params": G.g_synthesis.parameters()}, {"params": G.g_mapping.parameters(), "lr": lr_G * 0.01}] optim_G = torch.optim.Adam(params_G, lr=lr_G, betas=(0.9, 0.999)) scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99) scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99) # Train fix_z = torch.randn([opts.batch_size, 512]).to(opts.device) softplus = torch.nn.Softplus() Loss_D_list = [0.0] Loss_G_list = [0.0] for ep in range(start_epoch, opts.epoch): bar = tqdm(loader) loss_D_list = [] loss_G_list = [] for i, (real_img,) in enumerate(bar): real_img = real_img.to(opts.device) latents = torch.randn([real_img.size(0), 512]).to(opts.device) # ======================================================================================================= # (1) Update D network: D_logistic_r1(default) # ======================================================================================================= # Compute adversarial loss toward discriminator real_img = real_img.to(opts.device) real_logit = D(real_img) fake_img, fake_dlatent = G(latents) fake_logit = D(fake_img.detach()) d_loss = softplus(fake_logit) d_loss = d_loss + softplus(-real_logit) # original r1_penalty = D_logistic_r1(real_img.detach(), D) d_loss = (d_loss + r1_penalty).mean() # lite # d_loss = d_loss.mean() loss_D_list.append(d_loss.mean().item()) # Update discriminator optim_D.zero_grad() d_loss.backward() optim_D.step() # ======================================================================================================= # (2) Update G network: G_logistic_ns_pathreg(default) # ======================================================================================================= # if i % CRITIC_ITER == 0: G.zero_grad() fake_scores_out = D(fake_img) _g_loss = softplus(-fake_scores_out) # Compute |J*y|. # pl_noise = (torch.randn(fake_img.shape) / np.sqrt(fake_img.shape[2] * fake_img.shape[3])).to(fake_img.device) # pl_grads = grad(torch.sum(fake_img * pl_noise), fake_dlatent, retain_graph=True)[0] # pl_lengths = torch.sqrt(torch.sum(torch.sum(torch.mul(pl_grads, pl_grads), dim=2), dim=1)) # pl_mean = PL_DECAY * torch.sum(pl_lengths) # # pl_penalty = torch.mul(pl_lengths - pl_mean, pl_lengths - pl_mean) # reg = pl_penalty * PL_WEIGHT # # # original # g_loss = (_g_loss + reg).mean() # lite g_loss = _g_loss.mean() loss_G_list.append(g_loss.mean().item()) # Update generator g_loss.backward(retain_graph=True) optim_G.step() # Output training stats bar.set_description( "Epoch {} [{}, {}] [G]: {} [D]: {}".format(ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1])) # Save the result Loss_G_list.append(np.mean(loss_G_list)) Loss_D_list.append(np.mean(loss_D_list)) # Check how the generator is doing by saving G's output on fixed_noise with torch.no_grad(): fake_img = G(fix_z)[0].detach().cpu() save_image(fake_img, os.path.join(opts.det, 'images', str(ep) + '.png'), nrow=4, normalize=True) # Save model state = { 'G': G.state_dict(), 'D': D.state_dict(), 'Loss_G': Loss_G_list, 'Loss_D': Loss_D_list, 'start_epoch': ep, } torch.save(state, os.path.join(opts.det, 'models', 'latest.pth')) scheduler_D.step() scheduler_G.step() # Plot the total loss curve Loss_D_list = Loss_D_list[1:] Loss_G_list = Loss_G_list[1:] plotLossCurve(opts, Loss_D_list, Loss_G_list)
model.load_state_dict(torch.load(args.model_path)) # Prepare image and mask img = cv2.imread(args.image_path) origin_size = (np.shape(img)[1], np.shape(img)[0]) if args.mask_path is not None: mask = cv2.imread(args.mask_path) else: mask = generateMask(img) # Preprocessing proc_list = [ sunnertransforms.Rescale((args.size, args.size)), sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize() ] for op in proc_list: img = op(img) mask = op(mask) img = torch.stack([img, img], 0) mask = torch.stack([mask, mask], 0) mask = (mask + 1) / 2 # Work! if args.mask_path is not None: model.setInput(target=img, mask=mask) else: model.setInput(image=img, mask=mask) model.eval()
def main(opts): # Data load loader = sunnerData.DataLoader(sunnerData.ImageDataset( root=[[opts.path]], transform=transforms.Compose([ sunnertransforms.Resize((opts.resolution, opts.resolution)), sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize() ])), batch_size=opts.batch_size, shuffle=True, drop_last=True) # model generation start_epoch = 0 G = Generator_stylegan2(fmap_base=opts.fmap_base, resol=opts.resolution, mapping_layers=opts.mapping_layers, opts=opts, return_dlatents=True) D = Discriminator_stylegan2(fmap_base=opts.fmap_base, resol=opts.resolution, structure='resnet') # pre-trained weight loading if os.path.exists(opts.resume): INFO("Load the pre-trained weight!") state = torch.load(opts.resume) G.load_state_dict(state['G']) D.load_state_dict(state['D']) start_epoch = state['start_epoch'] else: INFO("pre-trained weight error") # multiple GPU support if (torch.cuda.device_count() > 1): INFO("multiple GPU detected! Total " + str(torch.cuda.device_count()) + '\t GPUs!') G = torch.nn.DataParrlel(G) D = torch.nn.DataParallel(D) G.to(opts.device) D.to(opts.device) # optimizer, scheduler lr_D = 0.0015 lr_G = 0.0015 optim_D = torch.optim.Adam(D.parameters(), lr=lr_D, betas=(0.9, 0.999)) params_G = [{ "params": G.g_synthesis.parameters() }, { "params": G.g_mapping.parameters(), "lr": lr_G * 0.01 }] optim_G = torch.optim.Adam(params_G, lr=lr_G, betas=(0.9, 0.999)) scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99) scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99) # start training fix_z = torch.randn([opts.batch_size, 512]).to(opts.device) softplus = torch.nn.Softplus() Loss_D_list = [0.0] Loss_G_list = [0.0] for ep in range(start_epoch, opts.epoch): bar = tqdm(loader) loss_D_list = [] loss_G_list = [] for i, (real_img, ) in enumerate(bar): real_img = real_img.to(opts.device) latents = torch.randn([real_img.size(0), 512]).to(opts.device) # Discriminator Network real_img = real_img.to(opts.device) real_logit = D(real_img) fake_img, fake_dlatent = G(latents) fake_logit = D(fake_img.detach()) d_loss = softplus(fake_logit) d_loss = d_loss + softplus(-real_logit) r1_penalty = D_logistic_r1(real_img.detach(), D) d_loss = (d_loss + r1_penalty).mean() loss_D_list.append(d_loss.mean().item()) optim_D.zero_grad() d_loss.backward() optim_D.step() # Generator Network G.zero_grad() fake_scores_out = D(fake_img) _g_loss = softplus(-fake_scores_out) g_loss = _g_loss.mean() loss_G_list.append(g_loss.mean().item()) g_loss.backward() optim_G.step() bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format( ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1])) # save result Loss_G_list.append(np.mean(loss_G_list)) Loss_D_list.append(np.mean(loss_D_list)) with torch.no_grad(): fake_img = G(fix_z)[0].detach().cpu() save_image(fake_img, os.path.join(opts.det, 'images', str(ep) + '.png'), nrow=4, normalize=True) # save model state = { 'G': G.state_dict(), 'D': D.state_dict(), 'Loss_G': Loss_G_list, 'Loss_D': Loss_D_list, 'start_epoch': ep, } torch.save(state, os.path.join(opts.det, 'models', 'latest.pth')) scheduler_D.step() scheduler_G.step() Loss_D_list = Loss_D_list[1:] Loss_G_list = Loss_G_list[1:] plotLossCurve(opts, Loss_D_list, Loss_G_list)
def train(opts): # Create the loader loader = sunnerData.DataLoader(dataset=BracketedDataset( root=opts.folder, crop_size=opts.crop_size, transform=transforms.Compose([ sunnertransforms.ToTensor(), sunnertransforms.ToFloat(), sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW), sunnertransforms.Normalize(), ])), batch_size=opts.batch_size, shuffle=True, num_workers=0) # Create the model model = DeepFuse(device=opts.device) criterion = MEF_SSIM_Loss().to(opts.device) optimizer = Adam(model.parameters(), lr=0.0001) # Load pre-train model if os.path.exists(opts.resume): state = torch.load(opts.resume) Loss_list = state['loss'] model.load_state_dict(state['model']) else: Loss_list = [] # Train bar = tqdm(range(opts.epoch)) for ep in bar: loss_list = [] for (patch1, patch2) in loader: # Extract the luminance and move to computation device patch1, patch2 = patch1.to(opts.device), patch2.to(opts.device) patch1_lum = patch1[:, 0:1] patch2_lum = patch2[:, 0:1] # Forward and compute loss model.setInput(patch1_lum, patch2_lum) y_f = model.forward() loss, y_hat = criterion(y_1=patch1_lum, y_2=patch2_lum, y_f=y_f) loss_list.append(loss.item()) bar.set_description("Epoch: %d Loss: %.6f" % (ep, loss_list[-1])) # Update the parameters optimizer.zero_grad() loss.backward() optimizer.step() Loss_list.append(np.mean(loss_list)) # Save the training image if ep % opts.record_epoch == 0: img = fusePostProcess(y_f, y_hat, patch1, patch2, single=False) cv2.imwrite(os.path.join(opts.det, 'image', str(ep) + ".png"), img[0, :, :, :]) # Save the training model if ep % (opts.epoch // 5) == 0: model_name = str(ep) + ".pth" else: model_name = "latest.pth" state = {'model': model.state_dict(), 'loss': Loss_list} torch.save(state, os.path.join(opts.det, 'model', model_name)) # Plot the loss curve plt.clf() plt.plot(Loss_list, '-') plt.title("loss curve") plt.savefig(os.path.join(opts.det, 'image', "curve.png"))