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 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 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')
