def generate_batch_translation_dataset(self, batch_size, truncation=0.4): # generate noise -> BxN(2)x128 noise_A = truncated_noise_sample(batch_size=batch_size, truncation=truncation) noise_B = truncated_noise_sample(batch_size=batch_size, truncation=truncation) n_A = torch.from_numpy(noise_A).unsqueeze(0) n_B = torch.from_numpy(noise_B).unsqueeze(0) n_AB = torch.cat([n_A, n_B], dim=0).to(self.device) n_AB = n_AB.view((N // 2) * batch_size, -1) # (B*2)x128 # generate class -> BxN(2)x1000 idx_list = range(len(self.class_list)) idx_ab = [random.sample(idx_list, 2) for i in range(batch_size)] # 2 pair idx class_ab = [[ self.class_list[idx_ab[i][0]], self.class_list[idx_ab[i][1]] ] for i in range(len(idx_ab))] class_ab = np.array(class_ab).transpose() class_A = one_hot_from_int(class_ab[0], batch_size=batch_size) # Bx1000 numpy array class_B = one_hot_from_int(class_ab[1], batch_size=batch_size) c_A = torch.from_numpy(class_A).unsqueeze(0) c_B = torch.from_numpy(class_B).unsqueeze(0) c_AB = torch.cat([c_A, c_B], dim=0).to(self.device) c_AB = c_AB.view((N // 2) * batch_size, -1) # (B*2)x1000 with torch.no_grad(): output = self.big_gan(n_AB[A + B + A + B], c_AB[A + B + B + A], truncation) # generate class idx idx_ab = torch.from_numpy(np.array(idx_ab).transpose()) # 2xB return output.view(N, batch_size, 3, output.size(2), output.size(3)), idx_ab.long()
def generate_batch_morphing_dataset(self, batch_size, interp_noise=True, interp_class=True, truncation=0.4): # generate noise -> BxN(3)x128 noise_A = truncated_noise_sample(batch_size=batch_size, truncation=truncation) noise_B = truncated_noise_sample( batch_size=batch_size, truncation=truncation) if interp_noise else 0 n_A = torch.from_numpy(noise_A).expand(N, batch_size, -1) n_B = torch.from_numpy(noise_B).expand( N, batch_size, -1) if interp_noise else n_A.clone() # generate class -> BxN(3)x1000 idx_list = range(len(self.class_list)) idx_ab = [random.sample(idx_list, 2) for i in range(batch_size)] # 2 pair idx class_ab = [[ self.class_list[idx_ab[i][0]], self.class_list[idx_ab[i][1]] ] for i in range(len(idx_ab))] class_ab = np.array(class_ab).transpose() class_A = one_hot_from_int(class_ab[0], batch_size=batch_size) # Bx1000 numpy array class_B = one_hot_from_int(class_ab[1], batch_size=batch_size) c_A = torch.from_numpy(class_A).expand(N, batch_size, -1) c_B = torch.from_numpy(class_B).expand( N, batch_size, -1) if interp_class else c_A.clone() # # sample index for intermediate image # idx_0 = random.randint(0, n_frames-i_frames) # idx_1 = idx_0 + i_frames - 1 # idx_i = random.randint(idx_0, idx_1) # sample index for intermediate image, including end points # interpolation -> BxN(3)xZ # t = torch.linspace(start=0, end=1.0, steps=n_frames) # t = t[[idx_0, idx_i, idx_1]].unsqueeze(1).unsqueeze(1).expand(3, batch_size, 1) a = torch.cat([ torch.zeros(1, batch_size, 1), torch.rand(1, batch_size, 1), torch.ones(1, batch_size, 1) ]) n = self.interp_func(n_A, n_B, a).to(self.device) c = self.interp_func(c_A, c_B, a).to(self.device) n = n.view(N * batch_size, -1) # (B*N(3))xZ c = c.view(N * batch_size, -1) with torch.no_grad(): output = self.big_gan(n, c, truncation) # generate class idx idx_ab = torch.from_numpy(np.array(idx_ab).transpose()) # 2xB return output.view(N, batch_size, 3, output.size(2), output.size(3)), a[1].squeeze(dim=1), idx_ab.long()
def generate_class_AB(self, class_A, class_B, batch_size=0): if batch_size == 0: self.class_A = one_hot_from_int(class_A) self.class_B = one_hot_from_int(class_B) else: # for translation dataset class_A = one_hot_from_int(class_A, batch_size=batch_size) class_B = one_hot_from_int(class_B, batch_size=batch_size) # generate class -> BxN(1)x1000 c_A = torch.from_numpy(class_A).unsqueeze(0) c_B = torch.from_numpy(class_B).unsqueeze(0) c_AB = torch.cat([c_A, c_B], dim=0).to(self.device) self.c_AB = c_AB.view(2 * batch_size, -1) # (B*2)x1000
def generate_images(num_images, class_index, model, truncation: float = 1.): print("[info] initializing input variables...") batch_size = min(10, num_images) class_vec = one_hot_from_int(class_index, batch_size=num_images) * .97 noise_vec = truncated_noise_sample(truncation=truncation, batch_size=num_images) noise_tensor = torch.tensor(noise_vec, dtype=torch.float) class_tensor = torch.tensor(class_vec, dtype=torch.float) print('[info] preparing model inference...') noise_tensor = noise_tensor.to(device) class_tensor = class_tensor.to(device) print( f'[info] generating {num_images} images of class index {class_index}...' ) images = [] for i in range(0, num_images, batch_size): torch.cuda.empty_cache() noise_feed = noise_tensor[i:i + batch_size] class_feed = class_tensor[i:i + batch_size] with torch.no_grad(): output = model(noise_feed, class_feed, truncation) output_cpu = output.cpu().data.numpy() for out in output_cpu: image = np.array(toimage(out)) images.append(image) print('[info] done.\n') torch.cuda.empty_cache() return images
def translate_sentiment(self, audio_sentiment, noise=0.1): # Step 1: Find the closest word match to the audio sentiment dists, idx = self.database.kneighbors(audio_sentiment) song_words = [self.words[cur_index[0]] for cur_index in idx] song_words = [x.split(',')[0].lower() for x in song_words] word_ids = [self.class_ids[cur_index] for cur_index in idx] class_vectors = [] noise_vectors = [] # Calculate the weighted nearest class neighbors for i in range(len(word_ids)): cur_class_vector = one_hot_from_int(word_ids[i], batch_size=len(word_ids[i])) # If k = 1 if len(word_ids[i]) == 1: cur_class_vector = cur_class_vector[0] else: similarities = np.exp(-dists[i] / self.temperature) similarities = similarities / np.sum(similarities) similarities = similarities.reshape((-1, 1)) cur_class_vector = np.sum(cur_class_vector * similarities, axis=0) class_vectors.append(cur_class_vector) noise_vectors.append( truncated_noise_sample(truncation=noise, batch_size=1)[0]) class_vectors = np.float32(class_vectors) noise_vectors = np.float32(noise_vectors) return class_vectors, noise_vectors, song_words
def forward(self, z, y): class_vector = one_hot_from_int(y, y.size(0)) class_vector = torch.from_numpy(class_vector).cuda() to_return = self.model(z, class_vector, 1.) return torch.nn.functional.interpolate( to_return, size=int(self.imagenet_preprocessing_str[-3:]), mode='bilinear') #, antialias = True)
def generate_t(self, class_A=0, class_B=1, class_t=0.5, truncation=0.4): self.class_vector_A = torch.from_numpy(one_hot_from_int(class_A)) self.class_vector_B = torch.from_numpy(one_hot_from_int(class_B)) # Interpolation class_vector_t = lerp(self.class_vector_A, self.class_vector_B, torch.tensor(class_t, dtype=torch.float32)) # GPU. class_vector_t = class_vector_t.to(self.device) # Generate images with torch.no_grad(): output = self.big_gan(self.noise_vector, class_vector_t, truncation=0.4) # save_as_images(output) return output
def set_output_class(self, class_id): if isinstance(class_id, int): self.v_class = torch.from_numpy(one_hot_from_int([class_id])).to( self.device) self.outclass = f"class{class_id}" elif isinstance(class_id, str): self.outclass = class_id.replace(" ", "_") self.v_class = torch.from_numpy(one_hot_from_names([class_id])).to( self.device) else: raise RuntimeError(f"Unknown class identifier {class_id}")
def get_sample_image(self): n = truncated_noise_sample(batch_size=self.batch_size, truncation=self.truncation) c_idx = np.random.choice(self.class_list, size=1) c = one_hot_from_int(c_idx, batch_size=self.batch_size) # np to torch n = torch.from_numpy(n).to(self.device) c = torch.from_numpy(c).to(self.device) with torch.no_grad(): output = self.big_gan(n, c, self.truncation) return output
def evaluate_inversion(args, inverted_net_path): # Load saved inverted net device = 'cuda:{}'.format( args.gpu_ids[0]) if len(args.gpu_ids) > 0 else 'cpu' ckpt_dict = torch.load(inverted_net_path, map_location=device) # Build model, load parameters model_args = ckpt_dict['model_args'] inverted_net = models.ResNet18(**model_args) inverted_net = nn.DataParallel(inverted_net, args.gpu_ids) inverted_net.load_state_dict(ckpt_dict['model_state']) import pdb pdb.set_trace() # Get test images (CelebA) initial_generated_image_dir = '/deep/group/sharonz/generator/z_test_images/' initial_generated_image_name = '058004_crop.jpg' initial_generated_image = util.get_image(initial_generated_image_dir, initial_generated_image_name) initial_generated_image = initial_generated_image / 255. intiial_generated_image = initial_generated_image.cuda() inverted_noise = inverted_net(initial_generated_image) if 'BigGAN' in args.model: class_vector = one_hot_from_int(207, batch_size=batch_size) class_vector = torch.from_numpy(class_vector) num_params = int(''.join(filter(str.isdigit, args.model))) generator = BigGAN.from_pretrained(f'biggan-deep-{num_params}') generator = generator.to(args.device) generated_image = generator.forward(inverted_noise, class_vector, args.truncation) # Get difference btw initial and subsequent generated image # Save both return
def __getitem__(self, idx): if idx > self.n_iters: raise StopIteration # Sample the space idx = np.random.randint(0, 1000, self.in_batch) class_vectors = one_hot_from_int(idx, batch_size=self.in_batch) noise_vectors = truncated_noise_sample(truncation=self.noise, batch_size=self.in_batch) class_vectors = torch.tensor(class_vectors).to(config['device']) noise_vectors = torch.tensor(noise_vectors).to(config['device']) with torch.no_grad(): output = self.gan_model(noise_vectors, class_vectors, self.noise) # Convert to PIL Image output = output.detach().cpu().numpy() output = np.uint8(np.clip(((output + 1) / 2.0) * 256, 0, 255)) output = output.transpose((0, 2, 3, 1)) images = [] # Pre-process each image to feed them into image sentiment analysis model for i in range(output.shape[0]): cur_img = Image.fromarray(output[i]) cur_img = self.transform_image(cur_img) images.append(cur_img) images = torch.stack(images).to(config['device']) # Feed-forward image sentiment analysis sentiment = self.sentiment_model(images) class_vectors = class_vectors.cpu().detach().numpy() noise_vectors = noise_vectors.cpu().detach().numpy() sentiment = sentiment.cpu().detach().numpy() return class_vectors, noise_vectors, sentiment
def test(args): # Get loader for z-test loader = get_loader(args, phase='test') batch_size = args.batch_size class_vector = None # TODO: make into function that takes in args.model and returns the pretrained model # and also consider whether it's class conditional and what kind of class conditional (how many classes) -> probably just imagenet now, actually maybe cifar-10 too # and also consider add truncation sampling as option too - this should return model, z_test noise vec, and class_vec (optionally) if args.ckpt_path and not args.use_pretrained: model, ckpt_info = ModelSaver.load_model(args.ckpt_path, args.gpu_ids) else: if 'BigGAN' in args.model: num_params = int(''.join(filter(str.isdigit, args.model))) if 'perturbation' in args.loss_fn: # Use custom BigGAN with Perturbation Net wrapper model = models.BigGANPerturbationNet.from_pretrained( f'biggan-deep-{num_params}') else: # Use pretrained BigGAN from package model = BigGAN.from_pretrained(f'biggan-deep-{num_params}') z_test = truncated_noise_sample(truncation=args.truncation, batch_size=batch_size) z_test = torch.from_numpy(z_test) # Get class conditional label # 981 is baseball player # 207 is golden retriever # TODO: Conditional generation only class_vector = one_hot_from_int(207, batch_size=batch_size) class_vector = torch.from_numpy(class_vector) elif 'WGAN-GP' in args.model: generator_path = "/deep/group/gen-eval/model-training/src/GAN_models/improved-wgan-pytorch/experiments/exp4_wgan_gp/generator.pt" model = torch.load(generator_path) z_test = torch.randn(batch_size, 128) elif 'BEGAN' in args.model: generator_path = "/deep/group/gen-eval/model-training/src/GAN_models/BEGAN-pytorch/trained_models/64/models/gen_97000.pth" model = models.BEGANGenerator() model.load_state_dict(torch.load(generator_path)) z_test = np.random.uniform(-1, 1, size=(batch_size, 64)) z_test = torch.FloatTensor(z_test) # Freeze model instead of using .eval() for param in model.parameters(): param.requires_grad = False # If using perturbation net, learn perturbation layers if 'perturbation' in args.loss_fn: trainable_params = [] for name, param in model.named_parameters(): if 'perturb' in name: param.requires_grad = True trainable_params.append(param) print(f'Number of trainable params: {len(trainable_params)}') model = nn.DataParallel(model, args.gpu_ids) model = model.to(args.device) # Loss functions if 'mse' in args.loss_fn: pixel_criterion = torch.nn.MSELoss().to(args.device) else: pixel_criterion = torch.nn.L1Loss().to(args.device) if 'perceptual' in args.loss_fn: # Combination pixel-perceptual loss - Sec 3.2. By default, uses pixel L1. perceptual_criterion = torch.nn.L1Loss().to(args.device) perceptual_loss_weight = args.perceptual_loss_weight vgg_feature_extractor = models.VGGFeatureExtractor().to(args.device) vgg_feature_extractor.eval() elif 'perturbation' in args.loss_fn: # Perturbation network R. By default, uses pixel L1. # Sec 3.3: http://ganpaint.io/Bau_et_al_Semantic_Photo_Manipulation_preprint.pdf reg_loss_weight = args.reg_loss_weight # z_loss_fn = util.get_loss_fn(args.loss_fn, args) max_z_test_loss = 100. # TODO: actually put max value possible here # Get logger, saver logger = TestLogger(args) # saver = ModelSaver(args) TODO: saver for perturbation network R print(f'Logs: {logger.log_dir}') print(f'Ckpts: {args.save_dir}') # Run z-test in batches logger.log_hparams(args) while not logger.is_finished_training(): logger.start_epoch() for _, z_test_target, mask in loader: logger.start_iter() if torch.cuda.is_available(): mask = mask.cuda() z_test = z_test.cuda() z_test_target = z_test_target.cuda() #class_vector = class_vector.cuda() masked_z_test_target = z_test_target * mask obscured_z_test_target = z_test_target * (1.0 - mask) if 'perturbation' in args.loss_fn: # With backprop on only trainable parameters in perturbation net params = trainable_params + [z_test.requires_grad_()] z_optimizer = util.get_optimizer(params, args) else: # With backprop on only the input z, run one step of z-test and get z-loss z_optimizer = util.get_optimizer([z_test.requires_grad_()], args) with torch.set_grad_enabled(True): if class_vector is not None: z_probs = model.forward(z_test, class_vector, args.truncation).float() z_probs = (z_probs + 1) / 2. else: z_probs = model.forward(z_test).float() # Calculate the masked loss using z-test vector masked_z_probs = z_probs * mask z_loss = torch.zeros(1, requires_grad=True).to(args.device) pixel_loss = torch.zeros(1, requires_grad=True).to(args.device) pixel_loss = pixel_criterion(masked_z_probs, masked_z_test_target) if 'perceptual' in args.loss_fn: z_probs_features = vgg_feature_extractor(masked_z_probs) z_test_features = vgg_feature_extractor( masked_z_test_target).detach() perceptual_loss = torch.zeros(1, requires_grad=True).to( args.device) perceptual_loss = perceptual_criterion( z_probs_features, z_test_features) z_loss = pixel_loss + perceptual_loss_weight * perceptual_loss elif 'perturbation' in args.loss_fn: reg_loss = torch.zeros(1, requires_grad=True).to(args.device) for name, param in model.named_parameters(): if 'perturb' in name: delta = param - 1 reg_loss += torch.pow(delta, 2).mean() #sum() z_loss = pixel_loss + reg_loss_weight * reg_loss else: z_loss = pixel_loss # Backprop on z-test vector z_loss.backward() z_optimizer.step() z_optimizer.zero_grad() # Compute the full loss (without mask) and obscured loss (loss only on masked region) # For logging and final evaluation (obscured loss is final MSE), so not in backprop loop full_z_loss = torch.zeros(1) full_pixel_loss = torch.zeros(1) full_pixel_loss = pixel_criterion(z_probs, z_test_target) #.mean() obscured_z_probs = z_probs * (1.0 - mask) obscured_z_loss = torch.zeros(1) obscured_pixel_loss = torch.zeros(1) obscured_pixel_loss = pixel_criterion( obscured_z_probs, obscured_z_test_target) #.mean() if 'perceptual' in args.loss_fn: # Full loss z_probs_full_features = vgg_feature_extractor(z_probs).detach() z_test_full_features = vgg_feature_extractor( z_test_target).detach() full_perceptual_loss = torch.zeros(1) full_perceptual_loss = perceptual_criterion( z_probs_full_features, z_test_full_features) full_z_loss = full_pixel_loss + perceptual_loss_weight * full_perceptual_loss # Obscured loss z_probs_obscured_features = vgg_feature_extractor( z_probs).detach() z_test_obscured_features = vgg_feature_extractor( z_test_target).detach() obscured_perceptual_loss = torch.zeros(1) obscured_perceptual_loss = perceptual_criterion( z_probs_obscured_features, z_test_obscured_features) obscured_z_loss = obscured_pixel_loss + perceptual_loss_weight * obscured_perceptual_loss elif 'perturbation' in args.loss_fn: full_z_loss = full_pixel_loss + reg_loss_weight * reg_loss obscured_z_loss = obscured_pixel_loss + reg_loss_weight * reg_loss else: full_z_loss = full_pixel_loss obscured_z_loss = obscured_pixel_loss """# TODO: z_loss is not always MSE anymore - figure out desired metric if z_loss < max_z_test_loss: # Save MSE on obscured region # TODO: z_loss is not always MSE anymore - figure out desired metric final_metrics = {'z-loss': z_loss.item(), 'obscured-z-loss': obscured_z_loss.item()} logger._log_scalars(final_metrics) print('Recall (z loss - non obscured loss - if MSE)', z_loss) print('Precision (MSE value on masked region)', obscured_z_loss) """ # Log both train and eval model settings, and visualize their outputs logger.log_status( masked_probs=masked_z_probs, masked_loss=z_loss, masked_test_target=masked_z_test_target, full_probs=z_probs, full_loss=full_z_loss, full_test_target=z_test_target, obscured_probs=obscured_z_probs, obscured_loss=obscured_z_loss, obscured_test_target=obscured_z_test_target, save_preds=args.save_preds, ) logger.end_iter() logger.end_epoch() # Last log after everything completes logger.log_status( masked_probs=masked_z_probs, masked_loss=z_loss, masked_test_target=masked_z_test_target, full_probs=z_probs, full_loss=full_z_loss, full_test_target=z_test_target, obscured_probs=obscured_z_probs, obscured_loss=obscured_z_loss, obscured_test_target=obscured_z_test_target, save_preds=args.save_preds, force_visualize=True, )
(classname, muembd, *xlim, *ylim, truncation, np.random.randint(10000)))) plt.show() #%% plt.figure(figsize=[4, 4]) plt.imshow(imgs[0]) plt.axis("image") plt.axis("off") plt.tight_layout() plt.show() #%% savedir = r"C:\Users\zhanq\OneDrive - Washington University in St. Louis\Generator_Testing\BigGAN256" truncation = 0.7 muembd = 0.1 classname = [1, 2] class_vector = one_hot_from_int([classname], batch_size=len(classname)) ebd_class = model.embeddings(torch.from_numpy(class_vector).cuda()) ebd_vecs = muembd * torch.randn( (10, ebd_class.shape[1] )).cuda() + ebd_class # add Gaussian perturbation around with torch.no_grad(): output = model.generator( torch.cat((torch.zeros_like(ebd_vecs), ebd_vecs), dim=1), truncation) imgs = convert_to_images(output.cpu()) figh = plt.figure(figsize=[25, 3]) gs = figh.add_gridspec(1, len(imgs)) # 1d interpolation for i, img in enumerate(imgs): plt.subplot(gs[i]) plt.imshow(img)
def train_inverted_net(args): # Start by training an external model on samples of G(z) -> z inversion model = util.get_invert_model(args) model = nn.DataParallel(model, args.gpu_ids) model = model.to(args.device) print(f'{args.invert_model} num params {count_parameters(model)}') generator = util.get_model(args) if generator is not None: generator = nn.DataParallel(generator, args.gpu_ids) generator = generator.to(args.device) print(f'{args.model} num params {count_parameters(generator)}') else: # Load saved pairings (ProGAN/StyleGAN) pairing_dir = '/deep/group/gen-eval/model-training/src/GAN_models/stylegan' pairing_path = f'{pairing_dir}/otavio_sampled_output/pairing.csv' pairings = pd.read_csv(pairing_path) num_pairings = len(pairings) noise_targets = pairings['noise'] image_inputs = pairings['image'] if 'BigGAN' in args.model: class_vector = one_hot_from_int(207, batch_size=args.batch_size) class_vector = torch.from_numpy(class_vector) class_vector = class_vector.cuda() # TODO: remove bc cant use gpu in laoder i don't think #loader = get_loader(args, phase='invert') #logger = TestLogger(args) #logger.log_hparams(args) criterion = torch.nn.MSELoss().to(args.device) optimizer = util.get_optimizer(model.parameters(), args) for i in range(args.num_invert_epochs): if generator is not None: noise_target = util.get_noise(args) image_input = generator.forward(noise_target).float() image_input = (image_input + 1.) / 2. else: # TODO: make into loader idx = i % num_pairings noise_target = np.load(f'{pairing_dir}/{noise_targets[idx]}') noise_target = torch.from_numpy(noise_target).float() print(f'noise target shape {noise_target.shape}') image_input = np.array( Image.open(f'{pairing_dir}/{image_inputs[idx]}')) image_input = torch.from_numpy(image_input / 255.) image_input = image_input.float().unsqueeze(0) image_input = image_input.permute(0, 3, 1, 2) noise_target = noise_target.cuda() image_input = image_input.cuda() with torch.set_grad_enabled(True): probs = model.forward(image_input) loss = torch.zeros(1, requires_grad=True).to(args.device) loss = criterion(probs, noise_target) print(f'iter {i}: loss = {loss}') loss.backward() optimizer.step() optimizer.zero_grad() if i % 1 == 0: corres_image_input = image_input.detach().cpu() corres_np = util.convert_image_from_tensor(corres_image_input) # Run check - saving image if 'BigGAN' in args.model: predicted_image = generator.forward(probs, class_vector, truncation).float() else: if generator is not None: predicted_image = generator.forward(probs).float() predicted_image = predicted_image.detach().cpu() predicted_image = (predicted_image + 1) / 2. predicted_np = util.convert_image_from_tensor( predicted_image) if len(predicted_np.shape) == 4: predicted_np = predicted_np[0] corres_np = corres_np[0] visuals = util.concat_images([predicted_np, corres_np]) visuals_pil = Image.fromarray(visuals) timestamp = datetime.now().strftime('%b%d_%H%M%S%f') visuals_image_dir = f'predicted_inversion_images/{args.model}' os.makedirs(visuals_image_dir, exist_ok=True) visuals_image_path = f'{visuals_image_dir}/{timestamp}_{i}.png' visuals_pil.save(visuals_image_path) print(f'Saved {visuals_image_path}') else: # Save noise vector - do forward separately in tf env probs = probs.detach().cpu().numpy() pred_noise_dir = f'predicted_inversion_noise/{args.model}' os.makedirs(pred_noise_dir, exist_ok=True) pred_noise_path = f'{pred_noise_dir}/{args.model}_noise_{i}.npy' np.save(pred_noise_path, probs) print(f'Saved {pred_noise_path}') if i % 1 == 0: corres_image_input = image_input.detach().cpu() corres_np = util.convert_image_from_tensor(corres_image_input) if len(corres_np.shape) == 4: corres_np = corres_np[0] corres_pil = Image.fromarray(corres_np) timestamp = datetime.now().strftime('%b%d_%H%M%S%f') corres_image_dir = f'generated_images/{args.model}' os.makedirs(corres_image_dir, exist_ok=True) corres_image_path = f'{corres_image_dir}/{timestamp}_{i}.png' corres_pil.save(corres_image_path) # saver = ModelSaver(args) global_step = args.num_invert_epochs ckpt_dict = { 'ckpt_info': { 'global_step': global_step }, 'model_name': model.module.__class__.__name__, 'model_args': model.module.args_dict(), 'model_state': model.to('cpu').state_dict(), 'optimizer': optimizer.state_dict(), } ckpt_dir = os.path.join(args.save_dir, f'{args.model}') os.makedirs(ckpt_dir, exist_ok=True) ckpt_path = os.path.join( ckpt_dir, f'{args.invert_model}_step_{global_step}.pth.tar') torch.save(ckpt_dict, ckpt_path) print(f'Saved model to {ckpt_path}') import pdb pdb.set_trace() return model
import os from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample, save_as_images, display_in_terminal, one_hot_from_int) # OPTIONAL: if you want to have more information on what's happening, activate the logger as follows import logging logging.basicConfig(level=logging.INFO) # Load pre-trained model tokenizer (vocabulary) model = BigGAN.from_pretrained('biggan-deep-512') # Prepare a input truncation = 0.4 # class_vector = one_hot_from_names(['soap bubble', 'coffee', 'mushroom'], batch_size=3) # pip install nltk class_vector = one_hot_from_int([445, 445, 445], batch_size=3) noise_vector = truncated_noise_sample(truncation=truncation, batch_size=3) # All in tensors noise_vector = torch.from_numpy(noise_vector) class_vector = torch.from_numpy(class_vector) # If you have a GPU, put everything on cuda noise_vector = noise_vector.to('cuda') class_vector = class_vector.to('cuda') model.to('cuda') # Generate an image with torch.no_grad(): output = model(noise_vector, class_vector, truncation)
def create_class_vector(x, device, batch_size: int = 1): class_vector = one_hot_from_int(x, batch_size) return torch.from_numpy(class_vector).to(device)
plt.hist(dist_mat[dist_mat != 0], bins=50) plt.title("Distribution of Distance Matrix of BigGAN-Deep256") plt.xlabel("L2 Distance") plt.savefig(join(homedir, "EmbedVectDistanceDist.png")) plt.show() #%% #%% from nltk.corpus import wordnet as wn import pickle, urllib ImageNet_Classname = pickle.load( urllib.request.urlopen( 'https://gist.githubusercontent.com/yrevar/6135f1bd8dcf2e0cc683/raw/d133d61a09d7e5a3b36b8c111a8dd5c4b5d560ee/imagenet1000_clsid_to_human.pkl' )) #%% Visualize Images for each class truncation = 0.7 class_vector = one_hot_from_int(list(range(1000)), batch_size=1000) sample_n = class_vector.shape[0] # noise_vector = np.zeros((sample_n, 128)) numbering = list(range(1000)) #%% noise_vec = truncated_noise_sample(batch_size=1, dim_z=128, truncation=0.7, seed=None) #%% batch = 10 csr = 0 csr_end = 0 while csr_end < sample_n: csr_end = min(csr + batch, sample_n) # imgs = BigGAN_render(class_vector[csr:csr_end, :], noise_vector[csr:csr_end, :], truncation)