def test_reconstruction(self, epoch, beta_kl, beta_mmd, beta_disc, batch_size=256, save_location=None): # todo remove betas self.eval() _, generator_test, _ = self.dataset.data_loaders(batch_size) with torch.no_grad(): for i, (x, m) in enumerate(generator_test): x = cuda_variable(x) m = cuda_variable(m) recon_batch = self.forward(x, m=m, beta_kl=beta_kl, beta_mmd=beta_mmd, beta_disc=beta_disc, num_samples=1, train=False)['samples'] if i == 0: n = min(x.size(0), 8) comparison = torch.cat( [x[:n], recon_batch.view(batch_size, 1, 28, 28)[:n]]) save_image(comparison.cpu(), save_location, nrow=n)
def test_variation_lines(self, num_elements_per_dim=20, save_location=None): self.eval() batch_size = 1 _, generator_test, _ = self.dataset.data_loaders(batch_size) generator_test_it = generator_test.__iter__() for _ in range(random.randint(0, len(generator_test))): x, _ = next(generator_test_it) trajectory = [] x = cuda_variable(x) with torch.no_grad(): z_star, _, _, _ = self.encoder_z.forward(x.detach()) alpha_it = self.features.alpha_iterator( max_num_dimensions=2, num_elements_per_dim=num_elements_per_dim) for alpha in alpha_it: style_vector = self.features.style_vector_from_alpha( alpha=alpha) logdet = cuda_variable(torch.zeros(batch_size)) z, _, _, _ = self.z_star_to_z.forward( z_star=z_star, logdet=logdet, style_vector=style_vector) x_pred, samples = self.decoder.forward(z) trajectory.append(x_pred) trajectory.append(x) trajectory = torch.cat( [t.view(batch_size, 1, 28, 28) for t in trajectory], 0) save_image(trajectory.cpu(), save_location, nrow=num_elements_per_dim)
def random_alpha(self, batch_size): rand_label = cuda_variable( torch.randint(low=0, high=self.num_values, size=(batch_size, )).long()) rand_alpha_sigmoid = cuda_variable( torch.rand(batch_size, self.num_style_tokens_per_label)) return rand_label, rand_alpha_sigmoid
def alpha_iterator(self, max_num_dimensions=None, num_elements_per_dim=10, offset=0): """ :param max_num_dimensions: Only for sigmoid features :param num_elements_per_dim: :param offset: :return: """ label_iterator = (cuda_variable(torch.Tensor([d]).long()) for d in range(self.num_values)) if max_num_dimensions is None: max_num_dimensions = self.num_style_tokens_per_label else: # label is always the first dimension max_num_dimensions = max_num_dimensions - 1 # todo to remove, investigate why this does not break # max_num_dimensions = max_num_dimensions g = itertools.product(np.arange(0., 1., 1 / num_elements_per_dim), repeat=max_num_dimensions) remaining_dimensions = self.num_style_tokens_per_label - max_num_dimensions begin = (0.5, ) * offset end = (0.5, ) * (remaining_dimensions - offset) alpha_sigmoid_iterator = (cuda_variable( torch.Tensor(begin + t + end)).unsqueeze(0) for t in g) return itertools.product(label_iterator, alpha_sigmoid_iterator)
def eval_batch(self, i, j): """ Get a batch of images in a range with their attributes. """ assert i < j batch_x = normalize_images(self.images[i:j].cuda()) batch_y = self.attributes[i:j].cuda() return cuda_variable(batch_x, volatile=True), cuda_variable(batch_y, volatile=True)
def forward(self, x): batch_size = x.size(0) x = x.view(batch_size, -1) context = self.context(x) # normalizing flow noise = cuda_variable(torch.randn(batch_size, self.z_dim)) logdet = cuda_variable(torch.zeros(batch_size)) z_samples, logdet, _ = self.inf((noise, logdet, context)) return z_samples, logdet, context, noise
def alpha_iterator(self, max_num_dimensions=None, num_elements_per_dim=10, offset=0): """ :param max_num_dimensions: :param num_elements_per_dim: :param offset: :return: alphas have a batch_size of 1 """ if max_num_dimensions is None: max_num_dimensions = self.alpha_dim else: max_num_dimensions = min(self.alpha_dim, max_num_dimensions) g = itertools.product(np.arange(0., 1., 1 / num_elements_per_dim), repeat=max_num_dimensions) remaining_dimensions = self.alpha_dim - max_num_dimensions begin = (0.5, ) * offset end = (0.5, ) * (remaining_dimensions - offset) alpha_gen = (cuda_variable(torch.Tensor(begin + t + end)).unsqueeze(0) for t in g) return alpha_gen
def alpha(self, x, m): # TODO only used for debug if isinstance(x, tuple) or isinstance(x, list): batch_size = x[0].size(0) else: batch_size = x.size(0) return cuda_variable(torch.zeros(batch_size))
def alpha_iterator(self, max_num_dimensions=None, num_elements_per_dim=None): # assert max_num_dimensions is None or max_num_dimensions == 1 assert num_elements_per_dim is None or max_num_dimensions <= self.num_values return (cuda_variable(torch.Tensor([d]).long()) for d in range(self.num_values))
def test_random_variations(self, num_variations=20, num_elements_per_dim=20, save_location=None): self.eval() batch_size = num_variations _, generator_test, _ = self.dataset.data_loaders(batch_size) generator_test_it = generator_test.__iter__() for _ in range(random.randint(1, len(generator_test))): x, m = next(generator_test_it) trajectory = [] x = cuda_variable(x) m = cuda_variable(m) with torch.no_grad(): z_star, _, _, _ = self.encoder_z.forward(x.detach()) z_star = Normal(torch.zeros_like(z_star), torch.ones_like(z_star)).sample() alpha = self.features.random_alpha(batch_size) # todo to remove # alpha = self.features.random_alpha(1) # alpha = alpha.repeat(batch_size, 1) # alpha = [alpha[0].repeat(batch_size), alpha[1].repeat(batch_size, 1)] style_vector = self.features.style_vector_from_alpha(alpha=alpha) logdet = cuda_variable(torch.zeros(batch_size)) z, _, _, _ = self.z_star_to_z.forward(z_star=z_star, logdet=logdet, style_vector=style_vector) x_pred, samples = self.decoder.forward(z) trajectory.append(x_pred) # trajectory.append(x) trajectory = torch.cat( [t.view(batch_size, 1, 28, 28) for t in trajectory], 0) save_image(trajectory.cpu(), save_location, nrow=num_elements_per_dim)
def train_batch(self, bs): """ Get a batch of random images with their attributes. """ # image IDs idx = torch.LongTensor(bs).random_(len(self.images)) # select images / attributes batch_x = normalize_images(self.images.index_select(0, idx).cuda()) batch_y = self.attributes.index_select(0, idx).cuda() # data augmentation if self.v_flip and np.random.rand() <= 0.5: batch_x = batch_x.index_select( 2, torch.arange(batch_x.size(2) - 1, -1, -1).long().cuda()) if self.h_flip and np.random.rand() <= 0.5: batch_x = batch_x.index_select( 3, torch.arange(batch_x.size(3) - 1, -1, -1).long().cuda()) return cuda_variable(batch_x), cuda_variable(batch_y)
def alpha_iterator(self, max_num_dimensions=None, num_elements_per_dim=10, offset=0): """ :param max_num_dimensions: :param num_elements_per_dim: :param offset: :return: alphas have a batch_size of 1 """ if max_num_dimensions is None: max_num_dimensions = self.num_style_tokens + 1 g = itertools.product(np.arange(0., 1., 1 / num_elements_per_dim), repeat=max_num_dimensions) remaining_dimensions = self.num_style_tokens + 1 - max_num_dimensions begin = (0.01, ) * offset end = (0.01, ) * (remaining_dimensions - offset) alpha_gen = (cuda_variable(torch.Tensor(begin + t + end)).unsqueeze(0) for t in g) probs_gen = (t / t.sum(1, keepdim=True) for t in alpha_gen) return probs_gen
def forward(self, x, m, train, beta_kl, beta_mmd, num_samples=1, beta_disc=1.): x, m = self.repeat_(x, m, num_samples=num_samples) z_star, logdet_z_star, context, noise = self.encoder_z(x) # compute alpha_from_input style_vector = self.features(x, m) # z knowing true tokens z, logdet_z, _, u_z = self.z_star_to_z.forward( z_star=z_star, logdet=logdet_z_star, style_vector=style_vector) # compute weights x_reconstruct, samples = self.decoder.forward(z=z, x=x, train=train) # KL on z_star zero = cuda_variable(torch.zeros(1)) logqz_star = utils.log_normal(noise, zero, zero).sum(1) - logdet_z_star logpz_star = utils.log_normal(z_star, zero, zero).sum(1) kl_star = logqz_star - logpz_star # free bits or beta? # kl = 0.1 * torch.max(kl, torch.ones_like(kl) * self.bits) # KL on z zero = cuda_variable(torch.zeros(1)) logqz = utils.log_normal(noise, zero, zero).sum(1) - logdet_z logpz = utils.log_normal(z, zero, zero).sum(1) kl = logqz - logpz # compute z (prior) prior_distribution = Normal(torch.zeros_like(z_star), torch.ones_like(z_star)) z_prior = prior_distribution.sample() z_star_prior = prior_distribution.sample() # mmd on z star mmd_z_star = mmd_reg(z_tilde=z_star, z=z_star_prior) # mmd on z mmd_z = mmd_reg(z_tilde=z, z=z_prior) ce = self.ce(value=x_reconstruct, target=x) disc_reg = self.disc_reg(z_star=z_star, style_vector=style_vector) loss = ( ce + beta_kl * kl_star # + beta_kl * kl + beta_mmd * mmd_z # + beta_mmd * mmd_z_star - beta_disc * disc_reg) loss = loss.mean() acc = self.accuracy(x_reconstruct, x) monitored_quantities = dict(loss=loss.item(), ce=ce.mean().item(), kl_star=kl_star.mean().item(), kl=kl.mean().item(), disc_reg=disc_reg.mean().item(), mmd_z=mmd_z.mean().item(), mmd_z_star=mmd_z_star.mean().item(), acc=acc * 100) return dict(loss=loss, monitored_quantities=monitored_quantities, samples=samples)
def test_visualization(self, num_elements_per_dim=20, num_curves=6, save_location=None): self.eval() batch_size = 1 _, generator_test, _ = self.dataset.data_loaders(batch_size) generator_test_it = generator_test.__iter__() for _ in range(random.randint(0, len(generator_test))): x, m = next(generator_test_it) z_trajectory = [] img_trajectory = [] x_original = cuda_variable(x) x = cuda_variable(x.repeat(num_curves, 1, 1, 1)) m = cuda_variable(m.repeat(num_curves)) with torch.no_grad(): print(self.features.alpha(x, m)) z_star, _, _, _ = self.encoder_z.forward(x.detach()) alpha_it = self.features.alpha_iterator( max_num_dimensions=2, num_elements_per_dim=num_elements_per_dim) # original style_vector = self.features.forward(x, m) logdet = cuda_variable(torch.zeros(batch_size)) z, _, _, _ = self.z_star_to_z.forward(z_star=z_star, logdet=logdet, style_vector=style_vector) for curve_index, coords in enumerate(z): original_line = torch.cat( [ coords, cuda_variable(torch.Tensor([curve_index])), cuda_variable(torch.Tensor([-1])), # -1 for original ], 0) z_trajectory.append(to_numpy(original_line[None, :])) # variation curves for alpha_index, alpha in enumerate(alpha_it): style_vector = self.features.style_vector_from_alpha( alpha=alpha) logdet = cuda_variable(torch.zeros(batch_size)) z, _, _, _ = self.z_star_to_z.forward( z_star=z_star, logdet=logdet, style_vector=style_vector) x_pred, samples = self.decoder.forward(z) img_trajectory.append(x_pred) for curve_index, coords in enumerate(z): line = torch.cat([ coords, cuda_variable(torch.Tensor([curve_index])), cuda_variable(torch.Tensor([alpha_index])) ], 0) z_trajectory.append(to_numpy(line[None, :])) img_trajectory = torch.cat( [t.view(num_curves, 1, 28, 28) for t in img_trajectory], 0) img_trajectory = torch.cat( [img_trajectory, x_original.view(1, 1, 28, 28)], 0) save_image(img_trajectory.cpu(), 'results/img_trajectories_same_x.png', nrow=num_curves) z_trajectory = np.concatenate(z_trajectory, axis=0) np.savetxt('results/trajectories_same_x.csv', z_trajectory, delimiter=',') batch_size = num_curves _, generator_test, _ = self.dataset.data_loaders(batch_size) generator_test_it = generator_test.__iter__() for _ in range(random.randint(0, len(generator_test))): x, m = next(generator_test_it) trajectory = [] img_trajectory = [] x = cuda_variable(x) m = cuda_variable(m) with torch.no_grad(): z_star, _, _, _ = self.encoder_z.forward(x.detach()) alpha_it = self.features.alpha_iterator( max_num_dimensions=2, num_elements_per_dim=num_elements_per_dim) # original style_vector = self.features.forward(x, m) logdet = cuda_variable(torch.zeros(batch_size)) z, _, _, _ = self.z_star_to_z.forward(z_star=z_star, logdet=logdet, style_vector=style_vector) for curve_index, coords in enumerate(z): original_line = torch.cat( [ coords, cuda_variable(torch.Tensor([curve_index])), cuda_variable(torch.Tensor([-1])), # -1 for original ], 0) trajectory.append(to_numpy(original_line[None, :])) for alpha_index, alpha in enumerate(alpha_it): style_vector = self.features.style_vector_from_alpha( alpha=alpha) logdet = cuda_variable(torch.zeros(batch_size)) z, _, _, _ = self.z_star_to_z.forward( z_star=z_star, logdet=logdet, style_vector=style_vector) x_pred, samples = self.decoder.forward(z) img_trajectory.append(x_pred) for curve_index, coords in enumerate(z): line = torch.cat([ coords, cuda_variable(torch.Tensor([curve_index])), cuda_variable(torch.Tensor([alpha_index])) ], 0) trajectory.append(to_numpy(line[None, :])) img_trajectory = torch.cat( [t.view(num_curves, 1, 28, 28) for t in img_trajectory], 0) img_trajectory = torch.cat( [img_trajectory, x.view(batch_size, 1, 28, 28)], 0) save_image(img_trajectory.cpu(), 'results/img_trajectories.png', nrow=num_curves) trajectory = np.concatenate(trajectory, axis=0) np.savetxt('results/trajectories.csv', trajectory, delimiter=',')
def preprocessing(self, *tensors): x, m = tensors return cuda_variable(x), cuda_variable(m)
def random_alpha(self, batch_size): return cuda_variable(torch.rand(batch_size, self.num_style_tokens))
def alpha_iterator(self, max_num_dimensions=None, num_elements_per_dim=10, offset=0): alpha_gen = (cuda_variable(torch.zeros(1)) for _ in range(1)) return alpha_gen
def style_vector_from_alpha(self, alpha): return cuda_variable(torch.zeros( (alpha.size(0), self.style_token_dim)))
def random_alpha(self, batch_size): probs = torch.rand(batch_size, self.num_style_tokens + 1) probs = probs / probs.sum(1, keepdim=True) return cuda_variable(probs)
def random_alpha(self, batch_size): r_a = torch.randint(low=0, high=self.num_values, size=(batch_size, )).long() return cuda_variable(r_a)