def log_prob(self, z, log_det_jacobians): ln_var_adj = self.ln_var * self.xp.ones([self.adj_size]) ln_var_x = self.ln_var * self.xp.ones([self.x_size]) log_det_jacobians[0] = log_det_jacobians[0] - F.log( self.xp.array([self.x_size], dtype=self.xp.float32)) log_det_jacobians[1] = log_det_jacobians[1] - F.log( self.xp.array([self.adj_size], dtype=self.xp.float32)) negative_log_likelihood_adj = F.average( F.sum(F.gaussian_nll(z[1], self.xp.zeros(self.adj_size, dtype=self.xp.float32), ln_var_adj, reduce="no"), axis=1) - log_det_jacobians[1]) negative_log_likelihood_x = F.average( F.sum(F.gaussian_nll(z[0], self.xp.zeros(self.x_size, dtype=self.xp.float32), ln_var_x, reduce="no"), axis=1) - log_det_jacobians[0]) negative_log_likelihood_adj /= self.adj_size negative_log_likelihood_x /= self.x_size if negative_log_likelihood_x.array < 0: log.warning("negative nll for x!") return [negative_log_likelihood_x, negative_log_likelihood_adj]
def lf(x): mu, ln_var = self.encode(x) mean_mu, mean_sigma = calculate_means(mu, ln_var) batchsize = len(mu.data) std_mu, std_ln_var = generate_std_params(mu) # reconstruction loss rec_loss = 0 kl_loss = 0 for l in six.moves.range(k): z = F.gaussian(mu, ln_var) rec_loss += F.bernoulli_nll(x, self.decode(z, sigmoid=False)) / (k * batchsize) kl_loss += -F.gaussian_nll(z, mu, ln_var) / (k * batchsize) kl_loss += F.gaussian_nll(z, std_mu, std_ln_var) / (k * batchsize) self.rec_loss = rec_loss self.kl_loss = kl_loss self.loss = self.rec_loss + C * self.kl_loss chainer.report( { 'rec_loss': rec_loss, 'kl': self.kl_loss, 'loss': self.loss, 'mu': mean_mu, 'sigma': mean_sigma, }, observer=self ) return self.loss
def check_invalid_option(self, xp): x = chainer.Variable(xp.asarray(self.x)) mean = chainer.Variable(xp.asarray(self.mean)) ln_var = chainer.Variable(xp.asarray(self.ln_var)) with self.assertRaises(ValueError): F.gaussian_nll(x, mean, ln_var, 'invalid_option')
def reconst(self, x, preprocessor=None, postprocessor=None, k=1): if preprocessor: x = preprocessor(x) x = np.reshape(x[:x.shape[0] // self.frame * self.frame], (-1, 1, self.frame, x.shape[1])) if self.gpu >= 0: x = cuda.to_gpu(x, device=self.gpu) xp = cuda.get_array_module(x) loss = xp.zeros(len(x)) # ここはメモリ節約のためにわざわざ分割計算している for i in np.arange(len(x)): x_batch = xp.reshape(x[i], (-1, 1, self.frame, self.dim)) mu_e, ln_var_e = self.network.encode(x_batch) mu_e.unchain_backward() ln_var_e.unchain_backward() for l in range(k): z = F.gaussian(mu_e, ln_var_e) z.unchain_backward() mu_d, ln_var_d = self.network.decode(z) mu_d.unchain_backward() ln_var_d.unchain_backward() loss[i] += (F.gaussian_nll(chainer.Variable(x_batch), mu_d, ln_var_d) / k).data loss = cuda.to_cpu(loss) if postprocessor: loss = postprocessor(loss) return loss
def backward(self, y_fake, mi, y_real): generator_loss = F.softmax_cross_entropy( y_fake, Variable(np.ones(y_fake.shape[0], dtype=np.int32))) discriminator_loss = F.softmax_cross_entropy( y_fake, Variable(np.zeros(y_fake.shape[0], dtype=np.int32))) discriminator_loss += F.softmax_cross_entropy( y_real, Variable(np.ones(y_real.shape[0], dtype=np.int32))) discriminator_loss /= 2 mi = mi.data # Mutual Information loss # Sample continuous codes to learn rotation, thickness, etc. c_continuous = np.asarray(rnd_continuous(mi.shape[0], mi.shape[1]), dtype=np.float32) # Continuous loss - Fix standard deviation to 1, i.e. log variance is 0 mi_continuous_ln_var = np.empty_like(mi, dtype=np.float32) mi_continuous_ln_var.fill(1) # mi_continuous_ln_var.fill(1e-6) continuous_loss = F.gaussian_nll(mi, Variable(c_continuous), Variable(mi_continuous_ln_var)) continuous_loss /= mi.shape[0] generator_loss += continuous_loss return {'gen': generator_loss, 'dis': discriminator_loss}
def nll_gaussian(preds, target, variance, add_const=False): """ Calc a negative log likelihood with gaussian distribution. Args: preds (numpy.ndarray or cupy.ndarray): """ if add_const: xp = backends.cuda.get_array_module(target) if isinstance(variance, chainer.Variable): variance = variance.array variance = xp.log(variance, dtype=target.dtype) F.gaussian_nll(preds, target, variance) neg_log_p = ((preds - target)**2 / (2 * variance)) const = 0.5 * np.log(2 * np.pi * variance) neg_log_p += const return neg_log_p.sum() / (target.size(0) * target.size(1))
def nll_gaussian(preds, target, variance, add_const=False): """ Calc a negative log likelihood with gaussian distribution. Args: preds (numpy.ndarray or cupy.ndarray): """ if add_const: xp = backends.cuda.get_array_module(target) if isinstance(variance, chainer.Variable): variance = variance.array variance = xp.log(variance, dtype=target.dtype) F.gaussian_nll(preds, target, variance) neg_log_p = ((preds - target) ** 2 / (2 * variance)) const = 0.5 * np.log(2 * np.pi * variance) neg_log_p += const return neg_log_p.sum() / (target.size(0) * target.size(1))
def free_energy(self, x): #return -(free energy) enc_mu, enc_log_sigma_2 = self.encode(x) kl = F.gaussian_kl_divergence(enc_mu, enc_log_sigma_2) z = F.gaussian(enc_mu, enc_log_sigma_2) dec_mu, dec_log_sigma_2 = self.decode(z) nll = F.gaussian_nll(x, dec_mu, dec_log_sigma_2) return nll + kl
def free_energy(self,x): #return -(free energy) enc_mu, enc_log_sigma_2 = self.encode(x) kl = F.gaussian_kl_divergence(enc_mu,enc_log_sigma_2) z = F.gaussian(enc_mu,enc_log_sigma_2) dec_mu, dec_log_sigma_2 = self.decode(z) nll = F.gaussian_nll(x,dec_mu,dec_log_sigma_2) return nll+kl
def logp(self): d = self.mu.shape[0] mean = np.broadcast_to(self.mu, (self.n_particle, d)) ln_var = np.broadcast_to(2 * np.log(self.sigma), (self.n_particle, d)) logp = -F.gaussian_nll(self.theta, mean, ln_var, reduce='no') logp = F.sum(logp, axis=1).reshape(-1) logp = F.broadcast_to(logp, (self.n_particle, self.n_particle)) return logp
def check_gaussian_nll(self, x, mean, ln_var): if self.wrap_x: x = chainer.Variable(x) if self.wrap_m: mean = chainer.Variable(mean) if self.wrap_v: ln_var = chainer.Variable(ln_var) actual = cuda.to_cpu(F.gaussian_nll(x, mean, ln_var, self.reduce).data) testing.assert_allclose(self.expect, actual)
def __call__(self, xs, x_tilda, x_fake, gpu=None, test=False): xp = np if gpu is None else cuda.cupy batch_size = xs.data.shape[0] # real h = F.relu(self.c0(xs)) h = F.relu(self.b_c1(self.c1(h), test=test)) h = F.relu(self.b_c2(self.c2(h), test=test)) h_xs = F.relu(self.b_c3(self.c3(h), test=test)) h_xs = F.reshape(h_xs, (h_xs.data.shape[0], h_xs.data.shape[1] * h_xs.data.shape[2] * h_xs.data.shape[3])) h = F.relu(self.b_l0(self.l0(h_xs), test=test)) y_xs = self.l1(h) # tilda h = F.relu(self.c0(x_tilda)) h = F.relu(self.b_c1(self.c1(h), test=test)) h = F.relu(self.b_c2(self.c2(h), test=test)) h_x_tilda = F.relu(self.b_c3(self.c3(h), test=test)) h_x_tilda = F.reshape( h_x_tilda, (h_x_tilda.data.shape[0], h_x_tilda.data.shape[1] * h_x_tilda.data.shape[2] * h_x_tilda.data.shape[3])) h = F.relu(self.b_l0(self.l0(h_x_tilda), test=test)) y_x_tilda = self.l1(h) # fake h = F.relu(self.c0(x_fake)) h = F.relu(self.b_c1(self.c1(h), test=test)) h = F.relu(self.b_c2(self.c2(h), test=test)) h = F.relu(self.b_c3(self.c3(h), test=test)) h = F.reshape(h, (h.data.shape[0], h.data.shape[1] * h.data.shape[2] * h.data.shape[3])) h = F.relu(self.b_l0(self.l0(h), test=test)) y_fake = self.l1(h) # calulate similarity loss zeros = Variable( xp.zeros((batch_size, h_xs.data.shape[1]), dtype=np.float32)) l_dis = F.gaussian_nll(h_xs, h_x_tilda, zeros) # calculate discrimation loss l_gan = F.softmax_cross_entropy( y_xs, Variable(xp.ones(batch_size, dtype=np.int32))) l_gan += F.softmax_cross_entropy( y_x_tilda, Variable(xp.zeros(batch_size, dtype=np.int32))) l_gan += F.softmax_cross_entropy( y_fake, Variable(xp.zeros(batch_size, dtype=np.int32))) print "P(y_xs) :", F.sum(F.softmax(y_xs), axis=0).data / batch_size print "P(y_x_tilda) :", F.sum(F.softmax(y_x_tilda), axis=0).data / batch_size print "P(y_fake) : ", F.sum(F.softmax(y_fake), axis=0).data / batch_size return l_dis, l_gan
def __call__(self, x): xp = self.encoder.xp x = Variable(xp.asarray(x)) zm, zv = self.encoder((x,)) z = F.gaussian(zm, zv) mean, ln_var = self.decoder((z,)) kl_loss = F.gaussian_kl_divergence(zm, zv) nll_loss = F.gaussian_nll(x, mean, ln_var) loss = kl_loss + nll_loss return loss
def __call__(self, x, l): mu, sigma = self.encoder(x) self.KL = F.gaussian_kl_divergence(mu, sigma) self.loss = Variable(np.array(0, dtype=np.float32)) for i in range(l): sample = F.gaussian(mu, sigma) m, s = self.decoder(sample) self.loss += F.gaussian_nll(x, m, s) self.loss = self.loss / l + self.KL self.loss = self.loss / len(x) return self.loss
def sample_u_i(self, zs, ui_prev, batch_size=0, prior_sample=False): if prior_sample == False: # sampling mu = self.l_u_mu(zs) ln_var = self.l_u_ln_var(zs) u_i = F.gaussian(mu, ln_var) # calculate kl logq = -gaussian_nll(u_i, mu, ln_var) # W = self.hoge.W need reshape # logp = -gaussian_nll(u_i, self.alphas * ui_prev, W*log_evar) log_evar = F.tile(self.log_evar, (u_i.data.shape[0], self.u_dims)) logp = -gaussian_nll(u_i, self.alphas * ui_prev, log_evar) batchsize = len(mu.data) kl_u_i = (logq - logp) / batchsize return u_i, kl_u_i else: # sampling log_evar = F.tile(self.log_evar, (batch_size, self.u_dims)) u_i = F.gaussian(self.alphas * ui_prev, log_evar) return u_i
def sample_u_1(self, zs, batch_size=0, prior_sample=False): xp = cuda.get_array_module(zs) if prior_sample == False: # sampling mu = self.l_u_mu(zs) ln_var = self.l_u_ln_var(zs) u_1 = F.gaussian(mu, ln_var) # calculate kl logq = -gaussian_nll(u_1, mu, ln_var) mu_cond = Variable(xp.zeros_like(u_1, dtype=xp.float32)) log_pvar = F.tile(self.log_pvar, (u_1.data.shape[0], self.u_dims)) logp = -gaussian_nll(u_1, mu_cond, log_pvar) batchsize = len(mu.data) kl_u_1 = (logq - logp) / batchsize return u_1, kl_u_1 else: # sampling log_pvar = F.tile(self.log_pvar, (batch_size, self.u_dims)) mus = xp.zeros((batch_size, self.u_dims), dtype=xp.float32) u_1 = F.gaussian(Variable(mus), log_pvar) return u_1
def sample_x_hat(self, zs, xs=[], nrep=1, calc_rec_loss=True): mu = self.l_x_mu(zs) ln_var = self.l_x_ln_var(zs) x_hat = [] for i in range(0, nrep): x_hat.append(F.gaussian(mu, ln_var)) x_hat = F.mean(F.vstack(x_hat), axis=0) if calc_rec_loss == True: batchsize = len(mu.data) rec_loss = gaussian_nll(xs, mu, ln_var) / batchsize return x_hat, rec_loss else: return x_hat
def __call__(self, *args, beta=1.0): assert len(args) >= 2 x = args[:-1] t = args[-1] mu_e, ln_var_e = self.predictor.encode(*x) batchsize = len(mu_e.data) rec_loss = 0 for l in six.moves.range(self.k): z = F.gaussian(mu_e, ln_var_e) mu_d, ln_var_d = self.predictor.decode(z) rec_loss += F.gaussian_nll(t, mu_d, ln_var_d) / (self.k * batchsize) kl_loss = beta * F.gaussian_kl_divergence(mu_e, ln_var_e) / batchsize self.loss = rec_loss + kl_loss reporter_module.report({'loss': self.loss}, self) return self.loss
def __call__(self, x): a, (x_mu, x_var), (z_mu, z_var), (a_mu, a_var) = self.net(x) if self.a is None: self.a = a self.theta_x = (x_mu, x_var) self.theta_z = (z_mu, z_var) self.theta_a = (a_mu, a_var) self.recon_loss = chainer.Variable(self.xp.array(0) .astype(x[0].data.dtype)) self.loss = chainer.Variable(self.xp.array(0) .astype(x[0].data.dtype)) return self.loss recon_loss = F.gaussian_nll(x, self.theta_x[0], self.theta_x[1]) kl_loss = F.gaussian_kl_divergence(self.theta_z[0], self.theta_z[1]) expect_loss = recon_loss + kl_loss # baseline eloss = expect_loss.data if self.base is None: self.base = eloss else: self.base = self.base_decay * self.base \ + (1-self.base_decay) * eloss # action loss act_loss = (eloss - self.base) \ * F.gaussian_nll(chainer.Variable(self.a.data), self.theta_a[0], self.theta_a[1]) self.theta_x = (x_mu, x_var) self.theta_z = (z_mu, z_var) self.theta_a = (a_mu, a_var) self.recon_loss = recon_loss self.loss = expect_loss return expect_loss + act_loss
def output_raw(self, test_filepath, fileloader=np.load, preprocessor=None): x = fileloader(test_filepath) if preprocessor: x = preprocessor(x) reshaped_x = np.reshape(x[:x.shape[0] // self.frame * self.frame], (-1, 1, self.frame, x.shape[1])) if self.gpu >= 0: reshaped_x = chainer.cuda.to_gpu(reshaped_x, device=self.gpu) mu_e, ln_var_e = self.network.encode(reshaped_x) z = F.gaussian(mu_e, ln_var_e) mu_d, ln_var_d = self.network.decode(z) xp = cuda.get_array_module(mu_d.data) mu_d = xp.reshape(mu_d.data, (-1, x.shape[1])) ln_var_d = xp.reshape(ln_var_d.data, (-1, x.shape[1])) loss = F.gaussian_nll(chainer.Variable(reshaped_x.reshape((-1, x.shape[1]))), chainer.Variable(mu_d), chainer.Variable(ln_var_d), reduce='no') return x, chainer.cuda.to_cpu(mu_d), chainer.cuda.to_cpu(ln_var_d), chainer.cuda.to_cpu(loss.data)
def calc_loss(self, x, t): h = self.head_model.feature(x) # print('head_space', h) mu, log_var = self.forward(x) self.neg_log_like_loss = F.gaussian_nll( t, mu, log_var) # returns the sum of nll's z = F.gaussian(mu, log_var) self.mean_abs_error = F.mean_absolute_error(t, z) chainer.report({'nll': self.neg_log_like_loss}, self) chainer.report({'mae': self.mean_abs_error}, self) chainer.report({'sigma': F.mean(F.sqrt(F.exp(log_var)))}, self) self.total_loss = 0.1*self.mean_abs_error + \ (self.neg_log_like_loss / len(x)) chainer.report({'loss': self.total_loss}, self) return self.total_loss
def test(self, test_filepath, fileloader=np.load, preprocessor=None, postprocessor=None, k=1): x = fileloader(test_filepath) if preprocessor: x = preprocessor(x) x = np.reshape(x[:x.shape[0] // self.frame * self.frame], (-1, 1, self.frame, x.shape[1])) if self.gpu >= 0: x = chainer.cuda.to_gpu(x, device=self.gpu) mu_e, ln_var_e = self.network.encode(x) batchsize = len(mu_e.data) loss = 0 for l in range(k): z = F.gaussian(mu_e, ln_var_e) mu_d, ln_var_d = self.network.decode(z) loss += (F.gaussian_nll(chainer.Variable(x), mu_d, ln_var_d, reduce='no') / k).data xp = cuda.get_array_module(loss) loss = xp.sum(loss, axis=(1, 2, 3)) loss = cuda.to_cpu(loss) if postprocessor: loss = postprocessor(loss) return loss
def train(self, x_context, y_context, x_target, y_target): xp = cuda.get_array_module(x_context) x_all = F.concat((x_context, x_target), axis=0) y_all = F.concat((y_context, y_target), axis=0) # Map x and y to z z_mu_all, z_ln_var_all = self.map_xy_to_z(x=x_all, y=y_all) z_mu_context, z_ln_var_context = self.map_xy_to_z(x=x_context, y=y_context) zs = F.gaussian(z_mu_all, z_ln_var_all) zs_rep = F.tile(zs, (x_target.data.shape[0], 1)) # decoder dec_mu, dec_ln_var = self.decoder(zs=zs_rep, x_star=x_target) # Loss = Log-likelihood (reconstruction loss) & KL-divergence rec_loss = gaussian_nll(y_target, mean=dec_mu, ln_var=dec_ln_var) kl = _gau_kl(p_mu=z_mu_context, p_ln_var=z_ln_var_context, q_mu=z_mu_all, q_ln_var=z_ln_var_all) return rec_loss, kl
def lf(x): mu, ln_var = self.encode(x) mean_mu, mean_sigma = calculate_means(mu, ln_var) batchsize = len(mu.data) # reconstruction loss rec_loss = 0 for l in six.moves.range(k): z = F.gaussian(mu, ln_var) mu_, ln_var_ = self.decode(z) rec_loss += F.gaussian_nll(x, mu_, ln_var_) / (k * batchsize) self.rec_loss = rec_loss kl = gaussian_kl_divergence(mu, ln_var) / batchsize self.loss = self.rec_loss + C * kl chainer.report( { 'rec_loss': rec_loss, 'loss': self.loss, 'kl': kl, 'mu': mean_mu, 'sigma': mean_sigma, }, observer=self) return self.loss
def lf(z_t, z_t_plus_1, action, done_label, reset=True): k = self.k output_dim = self.output_dim if reset: self.reset_state() output = self.fprop(F.concat((z_t, action))) if self.predict_done: coef, mu, ln_var, done = output else: coef, mu, ln_var = output coef = F.reshape(coef, (-1, output_dim, k)) coef = F.softmax(coef, axis=2) mu = F.reshape(mu, (-1, output_dim, k)) ln_var = F.reshape(ln_var, (-1, output_dim, k)) z_t_plus_1 = F.repeat(z_t_plus_1, k, 1).reshape(-1, output_dim, k) normals = F.sum( coef * F.exp(-F.gaussian_nll(z_t_plus_1, mu, ln_var, reduce='no')), axis=2) densities = F.sum(normals, axis=1) nll = -F.log(densities) loss = F.sum(nll) if self.predict_done: done_loss = F.sigmoid_cross_entropy(done.reshape(-1, 1), done_label, reduce="no") done_loss *= (1. + done_label.astype("float32") * 9.) done_loss = F.mean(done_loss) loss = loss + done_loss return loss
def main(): try: os.mkdir(args.snapshot_path) except: pass xp = np using_gpu = args.gpu_device >= 0 if using_gpu: cuda.get_device(args.gpu_device).use() xp = cupy num_bins_x = 2**args.num_bits_x image_size = (28, 28) images = chainer.datasets.mnist.get_mnist(withlabel=False)[0] images = 255.0 * np.asarray(images).reshape((-1, ) + image_size + (1, )) if args.num_channels != 1: images = np.broadcast_to( images, (images.shape[0], ) + image_size + (args.num_channels, )) images = preprocess(images, args.num_bits_x) x_mean = np.mean(images) x_var = np.var(images) dataset = glow.dataset.Dataset(images) iterator = glow.dataset.Iterator(dataset, batch_size=args.batch_size) print(tabulate([ ["#", len(dataset)], ["mean", x_mean], ["var", x_var], ])) hyperparams = Hyperparameters(args.snapshot_path) hyperparams.levels = args.levels hyperparams.depth_per_level = args.depth_per_level hyperparams.nn_hidden_channels = args.nn_hidden_channels hyperparams.image_size = image_size hyperparams.num_bits_x = args.num_bits_x hyperparams.lu_decomposition = args.lu_decomposition hyperparams.num_image_channels = args.num_channels hyperparams.save(args.snapshot_path) hyperparams.print() encoder = Glow(hyperparams, hdf5_path=args.snapshot_path) if using_gpu: encoder.to_gpu() optimizer = Optimizer(encoder) # Data dependent initialization if encoder.need_initialize: for batch_index, data_indices in enumerate(iterator): x = to_gpu(dataset[data_indices]) encoder.initialize_actnorm_weights( x, reduce_memory=args.reduce_memory) break current_training_step = 0 num_pixels = args.num_channels * hyperparams.image_size[0] * hyperparams.image_size[1] # Training loop for iteration in range(args.total_iteration): sum_loss = 0 sum_nll = 0 start_time = time.time() for batch_index, data_indices in enumerate(iterator): x = to_gpu(dataset[data_indices]) x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape) denom = math.log(2.0) * num_pixels factorized_z_distribution, logdet = encoder.forward_step( x, reduce_memory=args.reduce_memory) logdet -= math.log(num_bins_x) * num_pixels negative_log_likelihood = 0 for (zi, mean, ln_var) in factorized_z_distribution: negative_log_likelihood += cf.gaussian_nll(zi, mean, ln_var) loss = (negative_log_likelihood / args.batch_size - logdet) / denom encoder.cleargrads() loss.backward() optimizer.update(current_training_step) current_training_step += 1 sum_loss += float(loss.data) sum_nll += float(negative_log_likelihood.data) / args.batch_size printr( "Iteration {}: Batch {} / {} - loss: {:.8f} - nll: {:.8f} - log_det: {:.8f}". format( iteration + 1, batch_index + 1, len(iterator), float(loss.data), float(negative_log_likelihood.data) / args.batch_size / denom, float(logdet.data) / denom)) log_likelihood = -sum_nll / len(iterator) elapsed_time = time.time() - start_time print( "\033[2KIteration {} - loss: {:.5f} - log_likelihood: {:.5f} - elapsed_time: {:.3f} min". format(iteration + 1, sum_loss / len(iterator), log_likelihood, elapsed_time / 60)) encoder.save(args.snapshot_path)
def check_gaussian_nll(self, x_data, mean_data, ln_var_data): x = chainer.Variable(x_data) mean = chainer.Variable(mean_data) ln_var = chainer.Variable(ln_var_data) actual = cuda.to_cpu(F.gaussian_nll(x, mean, ln_var).data) testing.assert_allclose(self.expect, actual)
def check_gaussian_nll(self, x_data, mean_data, ln_var_data): x = chainer.Variable(x_data) mean = chainer.Variable(mean_data) ln_var = chainer.Variable(ln_var_data) actual = cuda.to_cpu(F.gaussian_nll(x, mean, ln_var, self.reduce).data) testing.assert_allclose(self.expect, actual)
def main(): try: os.mkdir(args.snapshot_path) except: pass xp = np using_gpu = args.gpu_device >= 0 if using_gpu: cuda.get_device(args.gpu_device).use() xp = cupy num_bins_x = 2**args.num_bits_x assert args.dataset_format in ["png", "npy"] # Get datasets: if True: files = Path(args.dataset_path).glob("*.{}".format( args.dataset_format)) if args.dataset_format == "png": images = [] for filepath in files: image = np.array(Image.open(filepath)).astype("float32") image = preprocess(image, args.num_bits_x) images.append(image) assert len(images) > 0 images = np.asanyarray(images) elif args.dataset_format == "npy": images = [] for filepath in files: array = np.load(filepath).astype("float32") # TODO: Preprocess here array = preprocess(array, args.num_bits_x) images.append(array) assert len(images) > 0 num_files = len(images) images = np.asanyarray(images) images = images.reshape((num_files * images.shape[1], ) + images.shape[2:]) else: raise NotImplementedError # Print dataset information if True: x_mean = np.mean(images) x_var = np.var(images) dataset = glow.dataset.Dataset(images) iterator = glow.dataset.Iterator(dataset, batch_size=args.batch_size) print( tabulate([ ["#", len(dataset)], ["mean", x_mean], ["var", x_var], ])) # Hyperparameters' info if True: hyperparams = Hyperparameters() hyperparams.levels = args.levels hyperparams.depth_per_level = args.depth_per_level hyperparams.nn_hidden_channels = args.nn_hidden_channels hyperparams.image_size = images.shape[2:] hyperparams.num_bits_x = args.num_bits_x hyperparams.lu_decomposition = args.lu_decomposition hyperparams.squeeze_factor = args.squeeze_factor hyperparams.save(args.snapshot_path) hyperparams.print() encoder = Glow(hyperparams, hdf5_path=args.snapshot_path) if using_gpu: encoder.to_gpu() optimizer = Optimizer(encoder) # Data dependent initialization if encoder.need_initialize: for batch_index, data_indices in enumerate(iterator): x = to_gpu(dataset[data_indices]) encoder.initialize_actnorm_weights(x) break current_training_step = 0 num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1] # Training loop for iteration in range(args.total_iteration): sum_loss = 0 sum_nll = 0 sum_kld = 0 start_time = time.time() for batch_index, data_indices in enumerate(iterator): x = to_gpu(dataset[data_indices]) x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape) denom = math.log(2.0) * num_pixels factorized_z_distribution, logdet = encoder.forward_step(x) logdet -= math.log(num_bins_x) * num_pixels kld = 0 negative_log_likelihood = 0 factor_z = [] for (zi, mean, ln_var) in factorized_z_distribution: negative_log_likelihood += cf.gaussian_nll(zi, mean, ln_var) if args.regularize_z: kld += cf.gaussian_kl_divergence(mean, ln_var) factor_z.append(zi.data.reshape(zi.shape[0], -1)) factor_z = xp.concatenate(factor_z, axis=1) negative_log_likelihood += cf.gaussian_nll( factor_z, xp.zeros(factor_z.shape, dtype='float32'), xp.zeros(factor_z.shape, dtype='float32')) loss = (negative_log_likelihood + kld) / args.batch_size - logdet loss = loss / denom encoder.cleargrads() loss.backward() optimizer.update(current_training_step) current_training_step += 1 sum_loss += _float(loss) sum_nll += _float(negative_log_likelihood) / args.batch_size sum_kld += _float(kld) / args.batch_size printr( "Iteration {}: Batch {} / {} - loss: {:.8f} - nll: {:.8f} - kld: {:.8f} - log_det: {:.8f}\n" .format( iteration + 1, batch_index + 1, len(iterator), _float(loss), _float(negative_log_likelihood) / args.batch_size / denom, _float(kld) / args.batch_size, _float(logdet) / denom)) if (batch_index + 1) % 100 == 0: encoder.save(args.snapshot_path) mean_log_likelihood = -sum_nll / len(iterator) mean_kld = sum_kld / len(iterator) elapsed_time = time.time() - start_time print( "\033[2KIteration {} - loss: {:.5f} - log_likelihood: {:.5f} - kld: {:.5f} - elapsed_time: {:.3f} min\n" .format(iteration + 1, sum_loss / len(iterator), mean_log_likelihood, mean_kld, elapsed_time / 60)) encoder.save(args.snapshot_path)
def main(): try: os.mkdir(args.ckpt) except: pass xp = np using_gpu = args.gpu_device >= 0 if using_gpu: cuda.get_device(args.gpu_device).use() xp = cupy hyperparams = Hyperparameters(args.snapshot_path) hyperparams.print() num_bins_x = 2.0**hyperparams.num_bits_x num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1] encoder = Glow(hyperparams, hdf5_path=args.snapshot_path) if using_gpu: encoder.to_gpu() # Load picture x = np.array(Image.open(args.img)).astype('float32') x = preprocess(x, hyperparams.num_bits_x) x = to_gpu(xp.expand_dims(x, axis=0)) x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape) if True: # Print this image info: b = xp.zeros((1, 3, 128, 128)) z, fw_ldt = encoder.forward_step(x, b) fw_ldt -= math.log(num_bins_x) * num_pixels logpZ = 0 ez = [] factor_z = [] for (zi, mean, ln_var) in z: factor_z.append(zi.data) logpZ += cf.gaussian_nll(zi, mean, ln_var) ez.append(zi.data.reshape(-1, )) ez = np.concatenate(ez) logpZ2 = cf.gaussian_nll(ez, xp.zeros(ez.shape), xp.zeros(ez.shape)).data print(fw_ldt, logpZ, logpZ2) with encoder.reverse() as decoder: rx, _ = decoder.reverse_step(factor_z) rx_img = make_uint8(rx.data[0], num_bins_x) rx_img = Image.fromarray(rx_img) rx_img.save(args.t + 'ori_revx.png') np.save(args.t + 'ori_z.npy', ez.get()) # Construct epsilon class eps(chainer.Chain): def __init__(self, shape, glow_encoder): super().__init__() self.encoder = glow_encoder with self.init_scope(): self.b = chainer.Parameter(initializers.Zero(), (1, 3, 128, 128)) self.m = chainer.Parameter(initializers.One(), (3, 8, 8)) def forward(self, x): # b = cf.tanh(self.b) * 0.5 b = self.b # Not sure if implementation is wrong m = cf.softplus(self.m) # m = cf.repeat(m, 8, axis=2) # m = cf.repeat(m, 8, axis=1) # m = cf.repeat(m, 16, axis=2) # m = cf.repeat(m, 16, axis=1) # b = b * m # cur_x = cf.add(x, b) # cur_x = cf.clip(cur_x, -0.5,0.5) z = [] zs, logdet = self.encoder.forward_step(x, b) for (zi, mean, ln_var) in zs: z.append(zi) z = merge_factorized_z(z) # return z, zs, logdet, cf.batch_l2_norm_squared(b), xp.tanh(self.b.data*1), cur_x, m return z, zs, logdet, xp.sum(xp.abs(b.data)), self.b.data * 1, m, x def save(self, path): filename = 'loss_model.hdf5' self.save_parameter(path, filename, self) def save_parameter(self, path, filename, params): tmp_filename = str(uuid.uuid4()) tmp_filepath = os.path.join(path, tmp_filename) save_hdf5(tmp_filepath, params) os.rename(tmp_filepath, os.path.join(path, filename)) epsilon = eps(x.shape, encoder) if using_gpu: epsilon.to_gpu() # optimizer = Optimizer(epsilon) optimizer = optimizers.Adam().setup(epsilon) # optimizer = optimizers.SGD().setup(epsilon) epsilon.b.update_rule.hyperparam.lr = 0.01 epsilon.m.update_rule.hyperparam.lr = 0.1 print('init finish') training_step = 0 z_s = [] b_s = [] loss_s = [] logpZ_s = [] logDet_s = [] m_s = [] j = 0 for iteration in range(args.total_iteration): epsilon.cleargrads() z, zs, fw_ldt, b_norm, b, m, cur_x = epsilon.forward(x) fw_ldt -= math.log(num_bins_x) * num_pixels logpZ1 = 0 factor_z = [] for (zi, mean, ln_var) in zs: factor_z.append(zi.data) logpZ1 += cf.gaussian_nll(zi, mean, ln_var) logpZ2 = cf.gaussian_nll(z, xp.zeros(z.shape), xp.zeros(z.shape)).data # logpZ2 = cf.gaussian_nll(z, np.mean(z), np.log(np.var(z))).data logpZ = (logpZ2 + logpZ1) / 2 loss = b_norm + (logpZ - fw_ldt) loss.backward() optimizer.update() training_step += 1 z_s.append(z.get()) b_s.append(cupy.asnumpy(b)) m_s.append(cupy.asnumpy(m.data)) loss_s.append(_float(loss)) logpZ_s.append(_float(logpZ)) logDet_s.append(_float(fw_ldt)) printr( "Iteration {}: loss: {:.6f} - b_norm: {:.6f} - logpZ: {:.6f} - logpZ1: {:.6f} - logpZ2: {:.6f} - log_det: {:.6f} - logpX: {:.6f}\n" .format(iteration + 1, _float(loss), _float(b_norm), _float(logpZ), _float(logpZ1), _float(logpZ2), _float(fw_ldt), _float(logpZ) - _float(fw_ldt))) if iteration % 100 == 99: np.save(args.ckpt + '/' + str(j) + 'z.npy', z_s) np.save(args.ckpt + '/' + str(j) + 'b.npy', b_s) np.save(args.ckpt + '/' + str(j) + 'loss.npy', loss_s) np.save(args.ckpt + '/' + str(j) + 'logpZ.npy', logpZ_s) np.save(args.ckpt + '/' + str(j) + 'logDet.npy', logDet_s) # cur_x = make_uint8(cur_x[0].data, num_bins_x) # np.save(args.ckpt + '/'+str(j)+'image.npy', cur_x) np.save(args.ckpt + '/' + str(j) + 'm.npy', m_s) with encoder.reverse() as decoder: rx, _ = decoder.reverse_step(factor_z) rx_img = make_uint8(rx.data[0], num_bins_x) np.save(args.ckpt + '/' + str(j) + 'res.npy', rx_img) z_s = [] b_s = [] loss_s = [] logpZ_s = [] logDet_s = [] m_s = [] j += 1 epsilon.save(args.ckpt)
def main(): xp = np using_gpu = args.gpu_device >= 0 if using_gpu: cuda.get_device(args.gpu_device).use() xp = cupy hyperparams = Hyperparameters(args.snapshot_path) hyperparams.print() num_bins_x = 2.0**hyperparams.num_bits_x num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1] # Get Dataset: if True: assert args.dataset_format in ["png", "npy"] files = Path(args.dataset_path).glob("*.{}".format( args.dataset_format)) if args.dataset_format == "png": images = [] for filepath in files: image = np.array(Image.open(filepath)).astype("float32") image = preprocess(image, hyperparams.num_bits_x) images.append(image) assert len(images) > 0 images = np.asanyarray(images) elif args.dataset_format == "npy": images = [] for filepath in files: array = np.load(filepath).astype("float32") array = preprocess(array, hyperparams.num_bits_x) images.append(array) assert len(images) > 0 num_files = len(images) images = np.asanyarray(images) images = images.reshape((num_files * images.shape[1], ) + images.shape[2:]) else: raise NotImplementedError dataset = glow.dataset.Dataset(images) iterator = glow.dataset.Iterator(dataset, batch_size=1) print(tabulate([["#image", len(dataset)]])) encoder = Glow(hyperparams, hdf5_path=args.snapshot_path) if using_gpu: encoder.to_gpu() ori_x = [] enc_z = [] rev_x = [] fw_logdet = [] logpZ = [] logpZ2 = [] i = 0 with chainer.no_backprop_mode() and encoder.reverse() as decoder: for data_indices in iterator: i += 1 x = to_gpu(dataset[data_indices]) # 1x3x64x64 x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape) x_img = make_uint8(x[0], num_bins_x) ori_x.append(x_img) factorized_z_distribution, fw_ldt = encoder.forward_step(x) fw_ldt -= math.log(num_bins_x) * num_pixels fw_logdet.append(cupy.asnumpy(fw_ldt.data)) factor_z = [] ez = [] nll = 0 for (zi, mean, ln_var) in factorized_z_distribution: nll += cf.gaussian_nll(zi, mean, ln_var) factor_z.append(zi.data) ez.append(zi.data.reshape(-1, )) ez = np.concatenate(ez) enc_z.append(ez.get()) logpZ.append(cupy.asnumpy(nll.data)) logpZ2.append( cupy.asnumpy( cf.gaussian_nll(ez, np.mean(ez), np.log(np.var(ez))).data)) rx, _ = decoder.reverse_step(factor_z) rx_img = make_uint8(rx.data[0], num_bins_x) rev_x.append(rx_img) if i % 100 == 0: np.save(str(i) + '/ori_x.npy', ori_x) fw_logdet = np.array(fw_logdet) np.save(str(i) + '/fw_logdet.npy', fw_logdet) np.save(str(i) + '/enc_z.npy', enc_z) logpZ = np.array(logpZ) np.save(str(i) + '/logpZ.npy', logpZ) logpZ2 = np.array(logpZ2) np.save(str(i) + '/logpZ2.npy', logpZ2) np.save(str(i) + '/rev_x.npy', rev_x) ori_x = [] enc_z = [] rev_x = [] fw_logdet = [] logpZ = [] logpZ2 = [] return
# Mutual Information loss mi_categorical, mi_continuous_mean = F.split_axis( mi, [n_categorical], 1) # Categorical loss categorical_loss = F.softmax_cross_entropy(mi_categorical, categories, use_cudnn=False) # Continuous loss - Fix standard deviation to 1, i.e. log variance is 0 mi_continuous_ln_var = xp.empty_like(mi_continuous_mean.data, dtype=xp.float32) mi_continuous_ln_var.fill(1) # mi_continuous_ln_var.fill(1e-6) continuous_loss = F.gaussian_nll(mi_continuous_mean, Variable(c_continuous), Variable(mi_continuous_ln_var)) continuous_loss /= batch_size generator_loss += categorical_loss generator_loss += continuous_loss # Backprop generator_optimizer.zero_grads() generator_loss.backward() generator_optimizer.update() discriminator_optimizer.zero_grads() discriminator_loss.backward() discriminator_optimizer.update()
def _posterior_nll(self, posterior, o, s): mu, ln_var = posterior(o) return F.gaussian_nll(x=s, mean=mu, ln_var=ln_var, reduce='mean')
def _transition_pll(self, transition, s_current, s_next): mu, ln_var = transition(s_current) return -F.gaussian_nll(x=s_next, mean=mu, ln_var=ln_var, reduce='mean')