def generate(config, writer, logger): config = config.opt config.distributed = False data_set = CreateDataLoader(config).load_data() model = create_model(config) visualizer = Visualizer(config) web_dir = os.path.join(config.results_dir, config.name, '%s_%s' % (config.phase, config.which_epoch)) webpage = html.HTML( web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (config.name, config.phase, config.which_epoch)) is_first = True average_tensor = utils.load_average_img(config) prev_generated = average_tensor.view(1, *average_tensor.shape) for data in tqdm(data_set): if config.no_temporal_smoothing: data['label'] = data['label'][:, :1] assert data['label'].shape[1] == 1 generated = model.inference(data['label'].cuda(), data['inst'].cuda()) visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], config.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0]))]) img_path = data['path'] visualizer.save_images(webpage, visuals, img_path) else: data['label'] = data['label'][:, :1] assert data['label'].shape[1] == 1 generated = model.inference(data['label'].cuda(), data['inst'].cuda(), prev_generated.cuda()) prev_generated = generated is_first = False visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], config.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0]))]) img_path = data['path'] visualizer.save_images(webpage, visuals, img_path)
def transfer(opt=test_opt): ''' Transfer source video to target video ''' iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() visualizer = Visualizer(opt) print(f'# testing images = {len(data_loader)}') # create website web_dir = os.path.join(opt.results_dir, opt.name, f'{opt.phase}_{opt.which_epoch}')) webpage = html.HTML(web_dir, f'Experiment = {opt.name}, Phase = {opt.phase}, Epoch = {opt.which_epoch}') model = create_model(opt) model = model.cuda() for data in tqdm(dataset): minibatch = 1 generated = model.inference(data['label'], data['inst']) visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0]))]) img_path = data['path'] visualizer.save_images(webpage, visuals, img_path) webpage.save() torch.cuda.empty_cache()
def gan_main(): opt = TestOptions().parse(save=False) opt.nThreads = 1 # test code only supports nThreads = 1 opt.batchSize = 1 # test code only supports batchSize = 1 opt.serial_batches = True # no shuffle opt.no_flip = True # no flip data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() model = create_model(opt) visualizer = Visualizer(opt) # create website web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch)) webpage = html.HTML( web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch)) # test for i, data in enumerate(dataset): if i >= opt.how_many: break generated = model.inference(data['label'], data['inst']) visuals = OrderedDict([ ('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0])) ]) img_path = data['path'] print('process image... %s' % img_path) visualizer.save_images(webpage, visuals, img_path) webpage.save()
def get_current_visuals(self): return OrderedDict([('input_label', util.tensor2label(self.input_label, self.opt.label_nc)), ('input_image', util.tensor2im(self.input_image)), ('real_image', util.tensor2im(self.real_image)), ('synthesized_image', util.tensor2im(self.fake_image))])
def generate_label_color(inputs): label_batch = [] for i in range(len(inputs)): label_batch.append(util.tensor2label(inputs[i], NC)) label_batch = np.array(label_batch) label_batch = label_batch * 2 - 1 input_label = torch.from_numpy(label_batch) return input_label
def get_current_visuals(self, getLabel=False): mask = self.mask if self.mask is not None: mask = np.transpose(self.mask[0].cpu().float().numpy(), (1,2,0)).astype(np.uint8) dict_list = [('fake_image', self.fake_image), ('mask', mask)] if getLabel: # only output label map if needed to save bandwidth label = util.tensor2label(self.net_input.data[0], self.opt.label_nc) dict_list += [('label', label)] return OrderedDict(dict_list)
def get_current_visuals(self, getLabel=False): mask = self.mask if self.mask is not None: mask = np.transpose(self.mask[0].cpu().float().numpy(), (1,2,0)).astype(np.uint8) dict_list = [('fake_image', self.fake_image), ('mask', mask)] if getLabel: # only output label map if needed to save bandwidth label = util.tensor2label(self.net_input.data[0], self.opt.label_nc) dict_list += [('label', label)] return OrderedDict(dict_list)
def single_generation_from_update(save_path, fname, features, checkpoints_dir, classname, black=True): ''' Generate decoded segmentation map from input features Args: save_path (str), save generated masks to path fname (str), save generated masks with fname features (numpy array): input features to be decoded checkpoints_dir (str), load VAE weights from path classname (str), label taxonomy defined by dataset with classname black (boolean), black is True for regular generation; black is False for debugging, thus the generated mask is not in the format for cGAN input ''' vae_opt = initialize_option(classname) vae_opt.checkpoints_dir = checkpoints_dir vae_util.mkdirs(save_path) if vae_opt.share_decoder and vae_opt.share_encoder: if vae_opt.separate_clothing_unrelated: from models.separate_clothing_encoder_models import create_model as vae_create_model else: print('Only supports separating clothing and clothing-irrelevant') raise NotImplementedError else: print('Only supports sharing encoder and decoder among all parts') raise NotImplementedError model = vae_create_model(vae_opt) generated = model.generate_from_random(torch.Tensor(features)) #.cuda()) if black: vae_util.save_image( vae_util.tensor2label_black(generated.data[0], vae_opt.output_nc, normalize=True), os.path.join(save_path, '%s.png' % (fname))) else: vae_util.save_image( vae_util.tensor2label(generated.data[0], vae_opt.output_nc, normalize=True), os.path.join(save_path, '%s.png' % (fname)))
def test_transfer(source_dir, run_name, temporal_smoothing=False, live_run_name=None): import src.config.test_opt as opt opt.name = run_name opt.dataroot = source_dir opt.temporal_smoothing = temporal_smoothing if device == torch.device('cpu'): opt.gpu_ids = [] else: os.environ['CUDA_VISIBLE_DEVICES'] = "0" opt.checkpoints_dir = os.path.join(dir_name, '../../checkpoints') opt.results_dir = os.path.join(dir_name, '../../results') data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() #print(opt.load_pretrain) model = create_model(opt) if live_run_name is not None: opt.name = live_run_name visualizer = Visualizer(opt) web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch)) webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch)) generated = None for data in tqdm(dataset): if temporal_smoothing: if generated is None: previous_frame = torch.zeros((1, 3, opt.loadSize, opt.loadSize)) generated = model.inference(data['label'], data['inst'], previous_frame) else: generated = model.inference(data['label'], data['inst'], generated) else: generated = model.inference(data['label'], data['inst']) visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0]))]) img_path = data['path'] visualizer.save_images(webpage, visuals, img_path) webpage.save() torch.cuda.empty_cache()
def get_current_visuals1(self): return OrderedDict([ ('input_label', util.tensor2label(self.input_label, self.opt.label_nc)), ('init_predict_label', util.tensor2label(self.init_predict_label, self.opt.label_nc)), ('input_label1', util.tensor2label(self.input_label1, self.opt.label_nc)), ('predict_label', util.tensor2label(self.predict_label, self.opt.label_nc)), ('ori_predict_label', util.tensor2label(self.ori_predict_label, self.opt.label_nc)), ('target_label', util.tensor2label(self.target_label, self.opt.label_nc)), ])
def visualize(data, generated, opt): if opt.model == 'pix2pixHD': visuals = OrderedDict([ ('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0])), ('real_image', util.tensor2im(data['image'][0])) ]) visualizer.display_current_results(visuals, epoch, total_steps) elif opt.model == 'pix2pixHDts': syn = generated[0].data[0] inputs = torch.cat((data['label'], data['next_label']), dim=3) targets = torch.cat((data['image'], data['next_image']), dim=3) visuals = OrderedDict([('input_label', util.tensor2im(inputs[0], normalize=False)), ('synthesized_image', util.tensor2im(syn)), ('real_image', util.tensor2im(targets[0]))]) if opt.face: #display face generator on tensorboard minx, miny, maxx, maxy = data['face_coords'][0] res_face = generated[2].data[0] syn_face = generated[1].data[0] preres = generated[3].data[0] visuals = OrderedDict([ ('input_label', util.tensor2im(inputs[0], normalize=False)), ('synthesized_image', util.tensor2im(syn)), ('synthesized_face', util.tensor2im(syn_face)), ('residual', util.tensor2im(res_face)), ('real_face', util.tensor2im(data['image'][0][:, miny:maxy, minx:maxx])), # ('pre_residual', util.tensor2im(preres)), # ('pre_residual_face', util.tensor2im(preres[:, miny:maxy, minx:maxx])), ('input_face', util.tensor2im(data['label'][0][:, miny:maxy, minx:maxx], normalize=False)), ('real_image', util.tensor2im(targets[0])) ]) visualizer.display_current_results(visuals, epoch, total_steps)
loss_D.backward() model.module.optimizer_D.step() #call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"]) ############## Display results and errors ########## ### print out errors if total_steps % opt.print_freq == print_delta: errors = {k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items()} t = (time.time() - iter_start_time) / opt.batchSize visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) ### display output images if save_fake: visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0])), ('real_image', util.tensor2im(data['image'][0]))]) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model if total_steps % opt.save_latest_freq == save_delta: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.module.save('latest') np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d') # end of epoch iter_end_time = time.time() print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
def train(): opt = TrainOptions().parse() if opt.debug: opt.display_freq = 1 opt.print_freq = 1 opt.nThreads = 1 ### initialize dataset data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) if opt.dataset_mode == 'pose': print('#training frames = %d' % dataset_size) else: print('#training videos = %d' % dataset_size) ### initialize models modelG, modelD, flowNet = create_model(opt) visualizer = Visualizer(opt) iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') ### if continue training, recover previous states if opt.continue_train: try: start_epoch, epoch_iter = np.loadtxt(iter_path, delimiter=',', dtype=int) except: start_epoch, epoch_iter = 1, 0 print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) if start_epoch > opt.niter: modelG.module.update_learning_rate(start_epoch - 1) modelD.module.update_learning_rate(start_epoch - 1) if (opt.n_scales_spatial > 1) and (opt.niter_fix_global != 0) and ( start_epoch > opt.niter_fix_global): modelG.module.update_fixed_params() if start_epoch > opt.niter_step: data_loader.dataset.update_training_batch( (start_epoch - 1) // opt.niter_step) modelG.module.update_training_batch( (start_epoch - 1) // opt.niter_step) else: start_epoch, epoch_iter = 1, 0 ### set parameters n_gpus = opt.n_gpus_gen // opt.batchSize # number of gpus used for generator for each batch tG, tD = opt.n_frames_G, opt.n_frames_D tDB = tD * opt.output_nc s_scales = opt.n_scales_spatial t_scales = opt.n_scales_temporal input_nc = 1 if opt.label_nc != 0 else opt.input_nc output_nc = opt.output_nc opt.print_freq = lcm(opt.print_freq, opt.batchSize) total_steps = (start_epoch - 1) * dataset_size + epoch_iter total_steps = total_steps // opt.print_freq * opt.print_freq ### real training starts here for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() for idx, data in enumerate(dataset, start=epoch_iter): if total_steps % opt.print_freq == 0: iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize # whether to collect output images save_fake = total_steps % opt.display_freq == 0 _, n_frames_total, height, width = data['B'].size( ) # n_frames_total = n_frames_load * n_loadings + tG - 1 n_frames_total = n_frames_total // opt.output_nc n_frames_load = opt.max_frames_per_gpu * n_gpus # number of total frames loaded into GPU at a time for each batch n_frames_load = min(n_frames_load, n_frames_total - tG + 1) t_len = n_frames_load + tG - 1 # number of loaded frames plus previous frames fake_B_last = None # the last generated frame from previous training batch (which becomes input to the next batch) real_B_all, fake_B_all, flow_ref_all, conf_ref_all = None, None, None, None # all real/generated frames so far real_B_skipped, fake_B_skipped = [None] * t_scales, [ None ] * t_scales # temporally subsampled frames flow_ref_skipped, conf_ref_skipped = [None] * t_scales, [ None ] * t_scales # temporally subsampled flows for i in range(0, n_frames_total - t_len + 1, n_frames_load): # 5D tensor: batchSize, # of frames, # of channels, height, width input_A = Variable( data['A'][:, i * input_nc:(i + t_len) * input_nc, ...]).view(-1, t_len, input_nc, height, width) input_B = Variable( data['B'][:, i * output_nc:(i + t_len) * output_nc, ...]).view(-1, t_len, output_nc, height, width) inst_A = Variable(data['inst'][:, i:i + t_len, ...]).view( -1, t_len, 1, height, width) if len(data['inst'].size()) > 2 else None ###################################### Forward Pass ########################## ####### generator fake_B, fake_B_raw, flow, weight, real_A, real_Bp, fake_B_last = modelG( input_A, input_B, inst_A, fake_B_last) if i == 0: fake_B_first = fake_B[ 0, 0] # the first generated image in this sequence real_B_prev, real_B = real_Bp[:, : -1], real_Bp[:, 1:] # the collection of previous and current real frames ####### discriminator ### individual frame discriminator flow_ref, conf_ref = flowNet( real_B, real_B_prev) # reference flows and confidences fake_B_prev = real_B_prev[:, 0: 1] if fake_B_last is None else fake_B_last[ 0][:, -1:] if fake_B.size()[1] > 1: fake_B_prev = torch.cat( [fake_B_prev, fake_B[:, :-1].detach()], dim=1) losses = modelD( 0, reshape([ real_B, fake_B, fake_B_raw, real_A, real_B_prev, fake_B_prev, flow, weight, flow_ref, conf_ref ])) losses = [ torch.mean(x) if x is not None else 0 for x in losses ] loss_dict = dict(zip(modelD.module.loss_names, losses)) ### temporal discriminator loss_dict_T = [] # get skipped frames for each temporal scale if t_scales > 0: real_B_all, real_B_skipped = get_skipped_frames( real_B_all, real_B, t_scales, tD) fake_B_all, fake_B_skipped = get_skipped_frames( fake_B_all, fake_B, t_scales, tD) flow_ref_all, conf_ref_all, flow_ref_skipped, conf_ref_skipped = get_skipped_flows( flowNet, flow_ref_all, conf_ref_all, real_B_skipped, flow_ref, conf_ref, t_scales, tD) # run discriminator for each temporal scale for s in range(t_scales): if real_B_skipped[s] is not None and real_B_skipped[ s].size()[1] == tD: losses = modelD(s + 1, [ real_B_skipped[s], fake_B_skipped[s], flow_ref_skipped[s], conf_ref_skipped[s] ]) losses = [ torch.mean(x) if not isinstance(x, int) else x for x in losses ] loss_dict_T.append( dict(zip(modelD.module.loss_names_T, losses))) # collect losses loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5 loss_G = loss_dict['G_GAN'] + loss_dict[ 'G_GAN_Feat'] + loss_dict['G_VGG'] loss_G += loss_dict['G_Warp'] + loss_dict[ 'F_Flow'] + loss_dict['F_Warp'] + loss_dict['W'] if opt.add_face_disc: loss_G += loss_dict['G_f_GAN'] + loss_dict['G_f_GAN_Feat'] loss_D += (loss_dict['D_f_fake'] + loss_dict['D_f_real']) * 0.5 # collect temporal losses loss_D_T = [] t_scales_act = min(t_scales, len(loss_dict_T)) for s in range(t_scales_act): loss_G += loss_dict_T[s]['G_T_GAN'] + loss_dict_T[s][ 'G_T_GAN_Feat'] + loss_dict_T[s]['G_T_Warp'] loss_D_T.append((loss_dict_T[s]['D_T_fake'] + loss_dict_T[s]['D_T_real']) * 0.5) ###################################### Backward Pass ################################# optimizer_G = modelG.module.optimizer_G optimizer_D = modelD.module.optimizer_D # update generator weights optimizer_G.zero_grad() loss_G.backward() optimizer_G.step() # update discriminator weights # individual frame discriminator optimizer_D.zero_grad() loss_D.backward() optimizer_D.step() # temporal discriminator for s in range(t_scales_act): optimizer_D_T = getattr(modelD.module, 'optimizer_D_T' + str(s)) optimizer_D_T.zero_grad() loss_D_T[s].backward() optimizer_D_T.step() if opt.debug: call([ "nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free" ]) ############## Display results and errors ########## ### print out errors if total_steps % opt.print_freq == 0: t = (time.time() - iter_start_time) / opt.print_freq errors = { k: v.data.item() if not isinstance(v, int) else v for k, v in loss_dict.items() } for s in range(len(loss_dict_T)): errors.update({ k + str(s): v.data.item() if not isinstance(v, int) else v for k, v in loss_dict_T[s].items() }) visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) ### display output images if save_fake: if opt.label_nc != 0: input_image = util.tensor2label(real_A[0, -1], opt.label_nc) elif opt.dataset_mode == 'pose': input_image = util.tensor2im(real_A[0, -1, :3], normalize=False) if real_A.size()[2] == 6: input_image2 = util.tensor2im(real_A[0, -1, 3:], normalize=False) input_image[input_image2 != 0] = input_image2[ input_image2 != 0] else: c = 3 if opt.input_nc == 3 else 1 input_image = util.tensor2im(real_A[0, -1, :c], normalize=False) if opt.use_instance: edges = util.tensor2im(real_A[0, -1, -1:, ...], normalize=False) input_image += edges[:, :, np.newaxis] if opt.add_face_disc: ys, ye, xs, xe = modelD.module.get_face_region(real_A[0, -1:]) if ys is not None: input_image[ys, xs:xe, :] = input_image[ ye, xs:xe, :] = input_image[ ys:ye, xs, :] = input_image[ys:ye, xe, :] = 255 visual_list = [ ('input_image', input_image), ('fake_image', util.tensor2im(fake_B[0, -1])), ('fake_first_image', util.tensor2im(fake_B_first)), ('fake_raw_image', util.tensor2im(fake_B_raw[0, -1])), ('real_image', util.tensor2im(real_B[0, -1])), ('flow_ref', util.tensor2flow(flow_ref[0, -1])), ('conf_ref', util.tensor2im(conf_ref[0, -1], normalize=False)) ] if flow is not None: visual_list += [('flow', util.tensor2flow(flow[0, -1])), ('weight', util.tensor2im(weight[0, -1], normalize=False))] visuals = OrderedDict(visual_list) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model if total_steps % opt.save_latest_freq == 0: visualizer.vis_print( 'saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) modelG.module.save('latest') modelD.module.save('latest') np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d') if epoch_iter > dataset_size - opt.batchSize: epoch_iter = 0 break # end of epoch iter_end_time = time.time() visualizer.vis_print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) ### save model for this epoch if epoch % opt.save_epoch_freq == 0: visualizer.vis_print( 'saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) modelG.module.save('latest') modelD.module.save('latest') modelG.module.save(epoch) modelD.module.save(epoch) np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d') ### linearly decay learning rate after certain iterations if epoch > opt.niter: modelG.module.update_learning_rate(epoch) modelD.module.update_learning_rate(epoch) ### gradually grow training sequence length if (epoch % opt.niter_step) == 0: data_loader.dataset.update_training_batch(epoch // opt.niter_step) modelG.module.update_training_batch(epoch // opt.niter_step) ### finetune all scales if (opt.n_scales_spatial > 1) and (opt.niter_fix_global != 0) and ( epoch == opt.niter_fix_global): modelG.module.update_fixed_params()
# test loss = torch.nn.L1Loss() losses = [] paths = [] for i, data in enumerate((dataset)): if i >= opt.how_many: break generated = model.fake_inference(data['image'], data['label'], data['inst'], pose=data['pose'], normal=data['normal'], depth=data['depth']) visuals = [('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('input_inst', util.tensor2label(data['inst'][0], opt.label_nc))] if opt.feat_pose: visuals += [('input_pose', util.tensor2label(data['pose'][0], opt.feat_pose_num_bins))] if opt.feat_normal: visuals += [('input_normal', util.tensor2im(data['normal'][0]))] if opt.feat_depth: visuals += [('input_depth', util.tensor2im(data['depth'][0]))] visuals += [ ('real_image', util.tensor2im(data['image'][0])), ('%s_%s_%s' % (opt.phase, opt.which_epoch, opt.experiment_name), util.tensor2im(generated.data[0])), ]
save_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch)) print('Doing %d frames' % len(dataset)) for i, data in enumerate(dataset): if i >= opt.how_many: break if data['change_seq']: model.fake_B_prev = None _, _, height, width = data['A'].size() A = Variable(data['A']).view(1, -1, input_nc, height, width) B = Variable(data['B']).view(1, -1, opt.output_nc, height, width) if len(data['B'].size()) > 2 else None inst = Variable(data['inst']).view(1, -1, 1, height, width) if len(data['inst'].size()) > 2 else None generated = model.inference(A, B, inst) if opt.label_nc != 0: real_A = util.tensor2label(generated[1], opt.label_nc) else: c = 3 if opt.input_nc == 3 else 1 real_A = util.tensor2im(generated[1][:c], normalize=False) visual_list = [('real_A', real_A), ('fake_B', util.tensor2im(generated[0].data[0]))] visuals = OrderedDict(visual_list) img_path = data['A_path'] print('process image... %s' % img_path) visualizer.save_images(save_dir, visuals, img_path)
# Get ground-truth output with torch.no_grad(): generated_noattack = model.inference(data['label'], data['inst'], data['image']) # Attack adv_image, perturb = model.attack(data['label'], data['inst'], data['image'], target=generated_noattack) # Get output from adversarial sample with torch.no_grad(): generated, adv_img = model.inference_attack( data['label'], data['inst'], data['image'], perturb) visuals = OrderedDict([ ('original_label', util.tensor2label(data['label'][0], opt.label_nc)), ('input_label', util.tensor2label(adv_img.data[0], opt.label_nc)), ('attacked_image', util.tensor2im(generated.data[0])), ('noattack', util.tensor2im(generated_noattack.data[0])) ]) img_path = data['path'] print('process image... %s' % img_path) visualizer.save_images(webpage, visuals, img_path) # Compute metrics l1_error += F.l1_loss(generated, generated_noattack) l2_error += F.mse_loss(generated, generated_noattack) l0_error += (generated - generated_noattack).norm(0) min_dist += (generated - generated_noattack).norm(float('-inf')) if F.mse_loss(generated, generated_noattack) > 0.05: n_dist += 1
k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items() } t = (time.time() - iter_start_time) / opt.batchSize visualizer.print_current_errors(epoch, train_epoch_iter, errors, t, mode='train') visualizer.plot_current_errors(errors, total_steps, mode='train') ### display output images if save_train: visuals = OrderedDict([ ('input_label', util.tensor2label(train_data['image'][0])), ('synthesized_image', util.tensor2im(generated.data[0], train_data['image'][0])), ('real_image', util.tensor2imreal(train_data['label'][0], train_data['image'][0])) ]) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model if total_steps % opt.save_latest_freq == 0: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.module.save('latest') np.savetxt(iter_path, (epoch, train_epoch_iter),
reconstructed_comb_label = reconstructed['comb_recon_label'] reconstructed_obj_label = reconstructed['obj_recon_label'] # generated = model.generate({ 'label_map': Variable(data['label'], volatile=True), 'mask_obj_in': Variable(data['mask_object_in'], volatile=True), 'mask_ctx_in': Variable(data['mask_context_in'], volatile=True), 'mask_obj_out': Variable(data['mask_object_out'], volatile=True), 'mask_out': Variable(data['mask_out'], volatile=True), 'mask_obj_inst': Variable(data['mask_object_inst'], volatile=True), 'cls': Variable(data['cls'], volatile=True), 'mask_in': Variable(data['mask_in'], volatile=True) }) generated_comb_label = generated['comb_pred_label'] generated_obj_label = generated['obj_pred_label'] visuals = OrderedDict([ ('image', util.tensor2im(data['image'][0])), ('gt_label', util.tensor2label(data['label'][0], opt.label_nc)), ('input_context', util.tensor2label(data['mask_context_in'][0], opt.label_nc)), ('mask_in', util.tensor2im(data['mask_in'][0])), ('reconstructed_comb_label', util.tensor2label(reconstructed_comb_label.data[0], opt.label_nc)), #('reconstructed_obj_label', util.tensor2im(reconstructed_obj_label.data[0])), ('generated_comb_label', util.tensor2label(generated_comb_label.data[0], opt.label_nc)), #('generated_obj_label', util.tensor2im(generated_obj_label.data[0]))]) label_path = data['label_path'] print('process image... %s' % label_path) visualizer.save_images(webpage, visuals, label_path) webpage.save()
data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() model = create_model(opt) visualizer = Visualizer(opt) # create website web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch)) webpage = html.HTML( web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch)) # test for i, data in enumerate(dataset): if i >= opt.how_many: break if opt.model == 'Bpgan_GAN': generated, latent_vector = model.inference(data['label']) elif opt.model == 'Bpgan_GAN_Q': generated, latent_vector = model.inference(data['label'], Q_type='Hard') visuals = OrderedDict([ ('input_label', util.tensor2label(data['label'][0], 0)), ('synthesized_image', util.tensor2im(generated.data[0])), ('real_image', util.tensor2im(data['image'][0])) ]) print(latent_vector.shape) img_path = data['path'] print('process image... %s' % img_path) visualizer.save_images(webpage, visuals, img_path) webpage.save()
for loss_name in loss_names: loss_mean_temp[loss_name] = loss_mean_temp[loss_name].item() / loss_count # errors = {k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items()} errors = {k: v for k, v in loss_mean_temp.items()} t = (time.time() - iter_start_time) / opt.batchSize visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) for loss_name in loss_names: loss_mean_temp[loss_name] = 0 loss_count = 0 ### display output images if save_fake: if opt.debug_mask_part == True: visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('input_ori_label', util.tensor2label(data['ori_label'][0], opt.label_nc)), ('transfer_label', util.tensor2label(transfer_label.data[0], opt.label_nc)), ('transfer_image', util.tensor2im(transfer_image.data[0])), ('reconstruct_image', util.tensor2im(reconstruct.data[0])), ('real_image', util.tensor2im(data['bg_image'][0])) # ('parsing_label', util.tensor2label(label_out.data[0], opt.label_nc)), # ('real_parsing_label', util.tensor2label(real_parsing_label.data[0], opt.label_nc)), # ('reconstruct_left_eye', util.tensor2im(left_eye_reconstruct.data[0])), # ('reconstruct_right_eye', util.tensor2im(right_eye_reconstruct.data[0])), # ('reconstruct_skin', util.tensor2im(skin_reconstruct.data[0])), # ('reconstruct_hair', util.tensor2im(hair_reconstruct.data[0])), # ('reconstruct_mouth', util.tensor2im(mouth_reconstruct.data[0])), # ('mask_lefteye', util.tensor2im(left_eye_real.data[0])) ]) else:
verbose=True) exit(0) minibatch = 1 if opt.engine: generated = run_trt_engine(opt.engine, minibatch, [data['label'], data['inst']]) elif opt.onnx: generated = run_onnx(opt.onnx, opt.data_type, minibatch, [data['label'], data['inst']]) else: generated = model.inference(data['label'], data['inst'], data['image']) img_path = data['path'] print('process image %d... %s' % (i, img_path)) orig_im = util.tensor2label(data['label'][0], opt.label_nc) mask = util.tensor2label(data['mask'][0], opt.label_nc) synthesized_im = util.tensor2im(generated.data[0]) masked_im = apply_mask(orig_im, synthesized_im, mask) enhanced = enhance_brightening(orig_im, masked_im) enhanced1 = enhance_brightening(orig_im, masked_im, factor=1.75) enhanced2 = enhance_brightening(orig_im, masked_im, factor=2) recolored = recolor_im(orig_im, mask, dst_im) recolored_from_synth = recolor_im(synthesized_im, mask, dst_im) # visuals = OrderedDict([('input_label', orig_im), # ('synthesized_image', synthesized_im), # ('masked_image', masked_im), # ('enhanced', enhanced) # ])
elif opt.data_type == 8: data['label'] = data['label'].uint8() data['inst'] = data['inst'].uint8() if opt.export_onnx: print("Exporting to ONNX: ", opt.export_onnx) assert opt.export_onnx.endswith( "onnx"), "Export model file should end with .onnx" torch.onnx.export(model, [data['label'], data['inst']], opt.export_onnx, verbose=True) exit(0) minibatch = 1 if opt.engine: generated = run_trt_engine(opt.engine, minibatch, [data['label'], data['inst']]) elif opt.onnx: generated = run_onnx(opt.onnx, opt.data_type, minibatch, [data['label'], data['inst']]) else: generated = model.inference(data['label'], data['inst'], data['image']) visuals = OrderedDict([ ('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0])) ]) img_path = data['path'] print('process image... %s' % img_path) visualizer.save_images(webpage, visuals, img_path) webpage.save()
############## Display results and errors ########## ### print out errors if total_steps % opt.print_freq == print_delta: errors = {k: v.data.item() if not isinstance(v, int) else v for k, v in loss_dict.items()} t = (time.time() - iter_start_time) / opt.print_freq visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) #call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"]) ### display output images if save_fake: ricker = util.tensor2im(data['image'][0]) genereted = util.tensor2im(generated.data[0]) psf = util.tensor2label(data['label'][0], opt.label_nc) diff_image = ricker - genereted print("Acc:", score(ricker, genereted)) visuals = OrderedDict([('PSF', psf), ('Genereted', genereted), ('Ricker', ricker), ('Difference', diff_image)]) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model """if total_steps % opt.save_latest_freq == save_delta: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.module.save('latest') np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d')"""
def main(): opt = TrainOptions().parse() iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') if opt.continue_train: try: start_epoch, epoch_iter = np.loadtxt(iter_path, delimiter=',', dtype=int) except: start_epoch, epoch_iter = 1, 0 print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) else: start_epoch, epoch_iter = 1, 0 opt.print_freq = lcm(opt.print_freq, opt.batchSize) if opt.debug: opt.display_freq = 1 opt.print_freq = 1 opt.niter = 1 opt.niter_decay = 0 opt.max_dataset_size = 10 data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) model = create_model(opt) visualizer = Visualizer(opt) if opt.fp16: from apex import amp model, [optimizer_G, optimizer_D ] = amp.initialize(model, [model.optimizer_G, model.optimizer_D], opt_level='O1') model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids) else: optimizer_G, optimizer_D = model.module.optimizer_G, model.module.optimizer_D total_steps = (start_epoch - 1) * dataset_size + epoch_iter display_delta = total_steps % opt.display_freq print_delta = total_steps % opt.print_freq save_delta = total_steps % opt.save_latest_freq for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() if epoch != start_epoch: epoch_iter = epoch_iter % dataset_size for i, data in enumerate(dataset, start=epoch_iter): if total_steps % opt.print_freq == print_delta: iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize # whether to collect output images save_fake = total_steps % opt.display_freq == display_delta ############## Forward Pass ###################### losses, generated = model(Variable(data['label']), Variable(data['inst']), Variable(data['image']), Variable(data['feat']), infer=save_fake) # sum per device losses losses = [ torch.mean(x) if not isinstance(x, int) else x for x in losses ] loss_dict = dict(zip(model.module.loss_names, losses)) # calculate final loss scalar loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5 loss_G = loss_dict['G_GAN'] + loss_dict.get( 'G_GAN_Feat', 0) + loss_dict.get('G_VGG', 0) ############### Backward Pass #################### # update generator weights optimizer_G.zero_grad() if opt.fp16: with amp.scale_loss(loss_G, optimizer_G) as scaled_loss: scaled_loss.backward() else: loss_G.backward() optimizer_G.step() # update discriminator weights optimizer_D.zero_grad() if opt.fp16: with amp.scale_loss(loss_D, optimizer_D) as scaled_loss: scaled_loss.backward() else: loss_D.backward() optimizer_D.step() ############## Display results and errors ########## ### print out errors if total_steps % opt.print_freq == print_delta: errors = { k: v.data.item() if not isinstance(v, int) else v for k, v in loss_dict.items() } t = (time.time() - iter_start_time) / opt.print_freq visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) #call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"]) ### display output images if save_fake: visuals = OrderedDict([ ('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0])), ('real_image', util.tensor2im(data['image'][0])) ]) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model if total_steps % opt.save_latest_freq == save_delta: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.module.save('latest') np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d') if epoch_iter >= dataset_size: break # end of epoch iter_end_time = time.time() print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) ### save model for this epoch if epoch % opt.save_epoch_freq == 0: print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) model.module.save('latest') model.module.save(epoch) np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d') ### instead of only training the local enhancer, train the entire network after certain iterations if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global): model.module.update_fixed_params() ### linearly decay learning rate after certain iterations if epoch > opt.niter: model.module.update_learning_rate()
def main(): iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) start_epoch, epoch_iter = 1, 0 total_steps = (start_epoch - 1) * dataset_size + epoch_iter display_delta = total_steps % opt.display_freq print_delta = total_steps % opt.print_freq save_delta = total_steps % opt.save_latest_freq model = create_model(opt) model = model.cuda() visualizer = Visualizer(opt) for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() if epoch != start_epoch: epoch_iter = epoch_iter % dataset_size for i, data in enumerate(dataset, start=epoch_iter): iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize # whether to collect output images save_fake = total_steps % opt.display_freq == display_delta ############## Forward Pass ###################### losses, generated = model(Variable(data['label']), Variable(data['inst']), Variable(data['image']), Variable(data['feat']), infer=save_fake) # sum per device losses losses = [torch.mean(x) if not isinstance(x, int) else x for x in losses] loss_dict = dict(zip(model.loss_names, losses)) # calculate final loss scalar loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5 loss_G = loss_dict['G_GAN'] + loss_dict.get('G_GAN_Feat', 0) + loss_dict.get('G_VGG', 0) ############### Backward Pass #################### # update generator weights model.optimizer_G.zero_grad() loss_G.backward() model.optimizer_G.step() # update discriminator weights model.optimizer_D.zero_grad() loss_D.backward() model.optimizer_D.step() ############## Display results and errors ########## ### print out errors if total_steps % opt.print_freq == print_delta: errors = {k: v.data if not isinstance(v, int) else v for k, v in loss_dict.items()} # CHANGE: removed [0] after v.data t = (time.time() - iter_start_time) / opt.batchSize visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) ### display output images if save_fake: visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0])), ('real_image', util.tensor2im(data['image'][0]))]) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model if total_steps % opt.save_latest_freq == save_delta: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.save('latest') np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d') if epoch_iter >= dataset_size: break # end of epoch print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) ### save model for this epoch if epoch % opt.save_epoch_freq == 0: print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) model.save('latest') model.save(epoch) np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d') ### instead of only training the local enhancer, train the entire network after certain iterations if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global): model.update_fixed_params() ### linearly decay learning rate after certain iterations if epoch > opt.niter: model.update_learning_rate() torch.cuda.empty_cache()
def train(opt): iter_path = os.path.join(opt.checkpoints_dir, 'iter.txt') if opt.continue_train: try: start_epoch, epoch_iter = np.loadtxt(iter_path, delimiter=',', dtype=int) except: start_epoch, epoch_iter = 1, 0 print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) else: start_epoch, epoch_iter = 1, 0 if opt.debug: opt.display_freq = 1 opt.print_freq = 1 opt.niter = 1 opt.niter_decay = 0 opt.max_dataset_size = 10 if opt.mode == 'meta-train': print("--picked k random images") opt.txtfile_img, opt.txtfile_label = create_k_txtfile_rand( opt.txtfile_img, opt.txtfile_label, opt.checkpoints_dir + "rand_k_img.txt", opt.checkpoints_dir + "rand_k_label.txt", opt.k) data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) model = create_model(opt) visualizer = Visualizer(opt) total_steps = (start_epoch - 1) * dataset_size + epoch_iter display_delta = total_steps % opt.display_freq print_delta = total_steps % opt.print_freq save_delta = total_steps % opt.save_latest_freq for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() if epoch != start_epoch: epoch_iter = epoch_iter % dataset_size for i, data in enumerate(dataset, start=epoch_iter): iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize # whether to collect output images save_fake = total_steps % opt.display_freq == display_delta ############## Forward Pass ###################### losses, generated = model(Variable(data['label']), Variable(data['inst']), Variable(data['image']), Variable(data['feat']), infer=save_fake) # sum per device losses losses = [ torch.mean(x) if not isinstance(x, int) else x for x in losses ] loss_dict = dict(zip(model.module.loss_names, losses)) # calculate final loss scalar loss_D = (loss_dict.get('D_fake', 0) + loss_dict.get('D_real', 0)) * 0.5 loss_G = loss_dict.get('G_GAN', 0) + loss_dict.get( 'G_GAN_Feat', 0) + loss_dict.get('G_VGG', 0) total_loss_dict = {'loss_D': loss_D, 'loss_G': loss_G} ############### Backward Pass #################### # update generator weights model.module.optimizer_G.zero_grad() loss_G.backward() model.module.optimizer_G.step() # update discriminator weights if type(loss_D) != float: model.module.optimizer_D.zero_grad() loss_D.backward() model.module.optimizer_D.step() #call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"]) ############## Display results and errors ########## ### print out errors if total_steps % opt.print_freq == print_delta: errors = { k: v.item() if not isinstance(v, int) else v for k, v in loss_dict.items() } t = (time.time() - iter_start_time) / opt.batchSize visualizer.print_current_errors(epoch, epoch_iter, errors, t, total_loss_dict) visualizer.plot_current_errors(errors, total_steps, total_loss_dict) ### display output images if save_fake: visuals = OrderedDict([ ('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0])), ('real_image', util.tensor2im(data['image'][0])) ]) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model if total_steps % opt.save_latest_freq == save_delta: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.module.save('latest') if opt.mode == 'test_checkpoints': model.module.save(str(epoch) + '_' + str(total_steps)) np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d') if epoch_iter >= dataset_size: break # end of epoch print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) ### save model for this epoch if epoch % opt.save_epoch_freq == 0: print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) model.module.save('latest') model.module.save(epoch) np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d') ### instead of only training the local enhancer, train the entire network after certain iterations if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global): model.module.update_fixed_params() ### linearly decay learning rate after certain iterations if epoch > opt.niter: model.module.update_learning_rate() if opt.mode != "meta-train": model.module.save('latest') exit_cuda() return model, loss_dict, total_loss_dict
def main(): opt = TrainOptions().parse() iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') if opt.continue_train: try: start_epoch, epoch_iter = np.loadtxt(iter_path, delimiter=',', dtype=int) except: start_epoch, epoch_iter = 1, 0 print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) else: start_epoch, epoch_iter = 1, 0 if opt.debug: opt.display_freq = 1 opt.print_freq = 1 opt.niter = 1 opt.niter_decay = 0 opt.max_dataset_size = 10 data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) model = create_model(opt) visualizer = Visualizer(opt) total_steps = (start_epoch - 1) * dataset_size + epoch_iter for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() if epoch != start_epoch: epoch_iter = epoch_iter % dataset_size for i, data in enumerate(dataset, start=epoch_iter): iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize # whether to collect output images save_fake = total_steps % opt.display_freq == 0 losses, real, label, generated, res, comp, up = model( Variable(data['label']), Variable(data['image']), Variable(data['ds']), infer=save_fake) # sum per device losses losses = [ torch.mean(x) if not isinstance(x, int) else x for x in losses ] loss_dict = dict(zip(model.module.loss_names, losses)) # calculate final loss scalar loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5 loss_G = loss_dict['G_GAN'] + loss_dict['G_GAN_Feat'] + loss_dict[ 'G_VGG'] + loss_dict['G_DIS'] + loss_dict['G_SSIM'] ############### Backward Pass #################### # update generator weights model.module.optimizer_G.zero_grad() loss_G.backward() model.module.optimizer_G.step() # update discriminator weights model.module.optimizer_D.zero_grad() loss_D.backward() model.module.optimizer_D.step() ############## Display results and errors ########## ### print out errors if total_steps % opt.print_freq == 0: errors = { k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items() } #errors['psp_loss'] = psp_train_loss.avg t = (time.time() - iter_start_time) / opt.batchSize visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) ### display output images if save_fake: #for i in range(opt.batchSize): i = 0 visuals = OrderedDict([ ('input_label', util.tensor2label(data['label'][i], opt.label_nc)), ('fine_image', util.tensor2im(res.data[i])), ('comp_image', util.tensor2im(comp.data[i])), ('up_image', util.tensor2im(up.data[i])), ('synthesized_image', util.tensor2im(generated.data[i])), ('real_image', util.tensor2im(data['image'][i])) ]) visualizer.display_current_results(visuals, epoch, epoch_iter) ### save latest model if total_steps % opt.save_latest_freq == 0: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.module.save('latest') np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d') # end of epoch print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) ### save model for this epoch if epoch % opt.save_epoch_freq == 0: print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) model.module.save('latest') model.module.save(epoch) np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d') ### instead of only training the local enhancer, train the entire network after certain iterations if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global): model.module.update_fixed_params() ### linearly decay learning rate after certain iterations if epoch > opt.niter: model.module.update_learning_rate()
label_encodings, num_labels = model.encode_features(Variable( data['input'])) # convert tensor to numpy label_encodings = label_encodings.data.cpu().numpy() for label in label_dict: fname_feature_dict['_'.join([fname, str(label)])] = \ label_encodings[:, label_dict[label]*opt.nz: (label_dict[label]+1)*opt.nz] fname_feature_dict['_'.join([fname, str(0)])] = label_encodings[:, -1 * opt.nz:] if sanity_check: # Sanity check the label_encodings generated = model.generate_from_random(label_encodings) util.save_image( util.tensor2label(data['input'][0], opt.output_nc, normalize=False), os.path.join(img_dir, 'input_label_%s.jpg' % (fname))) util.save_image( util.tensor2label(generated.data[0], opt.output_nc, normalize=True), os.path.join(img_dir, 'synthesized_label_%s.jpg' % (fname))) dir_path = ROOT_DIR + '/separate_vae/results/Lab/demo/' if not os.path.exists(dir_path): print(f"Creating directory {dir_path}") os.makedirs(dir_path) if demo: save_name = os.path.join(dir_path, '%s_shape_codes.p' % opt.phase) else: save_name = os.path.join(img_dir, '%s_shape_codes.p' % opt.phase) with open(save_name, 'wb') as writefile:
if opt.export_onnx: print("Exporting to ONNX: ", opt.export_onnx) assert opt.export_onnx.endswith( "onnx"), "Export model file should end with .onnx" torch.onnx.export(model, [data['label'], data['inst']], opt.export_onnx, verbose=True) exit(0) minibatch = 1 if opt.engine: generated = run_trt_engine(opt.engine, minibatch, [data['label'], data['inst']]) elif opt.onnx: generated = run_onnx(opt.onnx, opt.data_type, minibatch, [data['label'], data['inst']]) else: generated = model.inference(data['A'], data['B'], data['B2']) visuals = OrderedDict([('input_label', util.tensor2label(data['A'][0], 0)), ('real_image', util.tensor2im(data['A2'][0])), ('synthesized_image', util.tensor2im(generated.data[0])), ('B', util.tensor2label(data['B'][0], 0)), ('B2', util.tensor2im(data['B2'][0]))]) img_path = data['path'] img_path[0] = str(i) print('process image... %s' % img_path) visualizer.save_images(webpage, visuals, img_path) webpage.save()
if total_steps % opt.print_freq == print_delta: errors = { k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items() } errors['loss_G'] = loss_G errors['loss_D'] = loss_D errors['loss_SD'] = loss_SD t = (time.time() - iter_start_time) / opt.batchSize visualizer.print_current_errors(epoch, epoch_iter, errors, t) visualizer.plot_current_errors(errors, total_steps) ### display output images if save_fake: visuals = OrderedDict([ ('input_label', util.tensor2label(data['A'][0], 256)), ('real_image', util.tensor2im(data['A2'][0])), ('synthesized_image', util.tensor2im(generated.data[0])), ('B', util.tensor2label(data['B'][0], 256)), ('B2', util.tensor2im(data['B2'][0])) ]) visualizer.display_current_results2(visuals, epoch, total_steps) ### save latest model if total_steps % opt.save_latest_freq == save_delta: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.module.save('latest') np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d') if opt.use_iter_decay and total_steps > opt.niter_iter:
if opt.data_type == 16: data['label'] = data['label'].half() data['inst'] = data['inst'].half() elif opt.data_type == 8: data['label'] = data['label'].uint8() data['inst'] = data['inst'].uint8() if opt.export_onnx: print ("Exporting to ONNX: ", opt.export_onnx) assert opt.export_onnx.endswith("onnx"), "Export model file should end with .onnx" torch.onnx.export(model, [data['label'], data['inst']], opt.export_onnx, verbose=True) exit(0) minibatch = 1 if opt.engine: generated = run_trt_engine(opt.engine, minibatch, [data['label'], data['inst']]) elif opt.onnx: generated = run_onnx(opt.onnx, opt.data_type, minibatch, [data['label'], data['inst']]) else: generated = model.inference(data['label'], data['inst'], data['image']) visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0]))]) img_path = data['path'] print('process image... %s' % img_path) visualizer.save_images(webpage, visuals, img_path) print("----{}s seconds----".format(time.time() - start_time)) webpage.save()
ans.append(path) f_w = open(target_path, 'w') f_w.writelines(ans) opt = TrainOptions().parse() opt.phase = 'val' write_temp(opt, "temp") opt.phase = "temp" opt.serial_batches = True data_loader = CreatePoseConDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) visualizer = Visualizer(opt) total_steps = 0 # (start_epoch-1) * dataset_size + epoch_iter display_delta = total_steps % opt.display_freq print_delta = total_steps % opt.print_freq save_delta = total_steps % opt.save_latest_freq for i, data in enumerate(dataset): if (i % 100 == 0): print((i, dataset_size)) visuals = OrderedDict([('input_label', util.tensor2label(data['A'][0][3:6, :, :], 0)), ('real_image', util.tensor2im(data['A2'][0]))]) visualizer.display_current_results2(visuals, 0, i)