def save_images_and_clusters(images, clusters, epoch, shape=(10, 10)):
for i in range(10):
name_i = "results/3_gmvae/epoch_{}/cluster_{}.png".format(epoch, i)
images_i = images[clusters == i, :]
if images_i.shape[0] == 0:
continue
save_image_collections(images_i, name_i, shape=shape)
for t in range(test_iters):
test_x_batch = x_test[t * test_batch_size:(t + 1) *
test_batch_size]
test_y_batch = t_test[t * test_batch_size:(t + 1) *
test_batch_size]
test_lb = sess.run(lower_bound,
feed_dict={
x: test_x_batch,
y: test_y_batch,
n_particles: 1
})
test_ll = sess.run(is_log_likelihood,
feed_dict={
x: test_x_batch,
y: test_y_batch,
n_particles: 1
})
test_lbs.append(test_lb)
test_lls.append(test_ll)
time_test += time.time()
print('>>> TEST ({:.1f}s)'.format(time_test))
print('>> Test lower bound = {}'.format(np.mean(test_lbs)))
print('>> Test log likelihood (IS) = {}'.format(np.mean(test_lls)))
if epoch % save_freq == 0:
images = sess.run(x_gen)
name = os.path.join(result_path, "vae.epoch.{}.png".format(epoch))
save_image_collections(images, name)
# In[ ]:
def main():
# Load MNIST
data_path = os.path.join(conf.data_dir, 'mnist.pkl.gz')
x_train, t_train, x_valid, t_valid, x_test, t_test = \
dataset.load_mnist_realval(data_path)
x_train = np.vstack([x_train, x_valid])
x_test = np.random.binomial(1, x_test, size=x_test.shape)
x_dim = x_train.shape[1]
# Define model parameters
z_dim = 40
# Build the computation graph
n_particles = tf.placeholder(tf.int32, shape=[], name='n_particles')
x_input = tf.placeholder(tf.float32, shape=[None, x_dim], name='x')
x = tf.to_int32(tf.less(tf.random_uniform(tf.shape(x_input)), x_input))
n = tf.shape(x)[0]
def log_joint(observed):
model = vae(observed, x_dim, z_dim, n, n_particles)
log_pz, log_px_z = model.local_log_prob(['z', 'x'])
return log_pz + log_px_z
variational = q_net({'x': x}, x_dim, z_dim, n_particles)
qz_samples, log_qz = variational.query('z', outputs=True,
local_log_prob=True)
lower_bound = zs.variational.elbo(log_joint,
observed={'x': x},
latent={'z': [qz_samples, log_qz]},
axis=0)
cost = tf.reduce_mean(lower_bound.sgvb())
lower_bound = tf.reduce_mean(lower_bound)
# Importance sampling estimates of marginal log likelihood
is_log_likelihood = tf.reduce_mean(
zs.is_loglikelihood(log_joint, {'x': x},
{'z': [qz_samples, log_qz]}, axis=0))
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
infer_op = optimizer.minimize(cost)
# Generate images
n_gen = 100
x_mean = vae({}, x_dim, z_dim, n_gen).outputs('x_mean')
x_gen = tf.reshape(x_mean, [-1, 28, 28, 1])
# Define training/evaluation parameters
epochs = 3000
batch_size = 128
iters = x_train.shape[0] // batch_size
save_freq = 10
test_freq = 10
test_batch_size = 400
test_iters = x_test.shape[0] // test_batch_size
result_path = "results/vae"
# Run the inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
time_epoch = -time.time()
np.random.shuffle(x_train)
lbs = []
for t in range(iters):
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
_, lb = sess.run([infer_op, lower_bound],
feed_dict={x_input: x_batch,
n_particles: 1})
lbs.append(lb)
time_epoch += time.time()
print('Epoch {} ({:.1f}s): Lower bound = {}'.format(
epoch, time_epoch, np.mean(lbs)))
if epoch % test_freq == 0:
time_test = -time.time()
test_lbs = []
test_lls = []
for t in range(test_iters):
test_x_batch = x_test[t * test_batch_size:
(t + 1) * test_batch_size]
test_lb = sess.run(lower_bound,
feed_dict={x: test_x_batch,
n_particles: 1})
test_ll = sess.run(is_log_likelihood,
feed_dict={x: test_x_batch,
n_particles: 1000})
test_lbs.append(test_lb)
test_lls.append(test_ll)
time_test += time.time()
print('>>> TEST ({:.1f}s)'.format(time_test))
print('>> Test lower bound = {}'.format(np.mean(test_lbs)))
print('>> Test log likelihood (IS) = {}'.format(
np.mean(test_lls)))
if epoch % save_freq == 0:
images = sess.run(x_gen)
name = os.path.join(result_path,
"vae.epoch.{}.png".format(epoch))
save_image_collections(images, name)
def main():
# Load MNIST
data_path = os.path.join(conf.data_dir, "mnist.pkl.gz")
x_train, t_train, x_valid, t_valid, x_test, t_test = \
dataset.load_mnist_realval(data_path)
x_train = np.vstack([x_train, x_valid])
x_test = np.random.binomial(1, x_test, size=x_test.shape)
x_dim = x_train.shape[1]
# Define model parameters
z_dim = 40
# Build the computation graph
n_particles = tf.placeholder(tf.int32, shape=[], name="n_particles")
x_input = tf.placeholder(tf.float32, shape=[None, x_dim], name="x")
x = tf.cast(tf.less(tf.random_uniform(tf.shape(x_input)), x_input),
tf.int32)
n = tf.placeholder(tf.int32, shape=[], name="n")
model = build_gen(x_dim, z_dim, n, n_particles)
variational = build_q_net(x, z_dim, n_particles)
lower_bound = zs.variational.elbo(model, {"x": x},
variational=variational,
axis=0)
cost = tf.reduce_mean(lower_bound.sgvb())
lower_bound = tf.reduce_mean(lower_bound)
# # Importance sampling estimates of marginal log likelihood
is_log_likelihood = tf.reduce_mean(
zs.is_loglikelihood(model, {"x": x}, proposal=variational, axis=0))
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
infer_op = optimizer.minimize(cost)
# Random generation
x_gen = tf.reshape(model.observe()["x_mean"], [-1, 28, 28, 1])
# Define training/evaluation parameters
epochs = 3000
batch_size = 128
iters = x_train.shape[0] // batch_size
save_freq = 10
test_freq = 10
test_batch_size = 400
test_iters = x_test.shape[0] // test_batch_size
result_path = "results/vae"
# Run the inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
time_epoch = -time.time()
np.random.shuffle(x_train)
lbs = []
for t in range(iters):
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
_, lb = sess.run([infer_op, lower_bound],
feed_dict={
x_input: x_batch,
n_particles: 1,
n: batch_size
})
lbs.append(lb)
time_epoch += time.time()
print("Epoch {} ({:.1f}s): Lower bound = {}".format(
epoch, time_epoch, np.mean(lbs)))
if epoch % test_freq == 0:
time_test = -time.time()
test_lbs, test_lls = [], []
for t in range(test_iters):
test_x_batch = x_test[t * test_batch_size:(t + 1) *
test_batch_size]
test_lb = sess.run(lower_bound,
feed_dict={
x: test_x_batch,
n_particles: 1,
n: test_batch_size
})
test_ll = sess.run(is_log_likelihood,
feed_dict={
x: test_x_batch,
n_particles: 1000,
n: test_batch_size
})
test_lbs.append(test_lb)
test_lls.append(test_ll)
time_test += time.time()
print(">>> TEST ({:.1f}s)".format(time_test))
print(">> Test lower bound = {}".format(np.mean(test_lbs)))
print('>> Test log likelihood (IS) = {}'.format(
np.mean(test_lls)))
if epoch % save_freq == 0:
images = sess.run(x_gen, feed_dict={n: 100, n_particles: 1})
name = os.path.join(result_path,
"vae.epoch.{}.png".format(epoch))
save_image_collections(images, name)
is_training: True
})
gen_losses.append(g_loss)
disc_losses.append(d_loss)
if iter % print_freq == 0:
print('Epoch={} Iter={} ({:.3f}s/iter): '
'Gen loss = {} Disc loss = {}'.format(
epoch, iter,
(time.time() + time_train) / print_freq,
np.mean(gen_losses), np.mean(disc_losses)))
gen_losses = []
disc_losses = []
if iter % test_freq == 0:
time_test = -time.time()
images = sess.run(eval_x_gen,
feed_dict={is_training: False})
name = "results/dcgan/dcgan.epoch.{}.iter.{}.png".format(
epoch, iter)
save_image_collections(images, name, scale_each=True)
time_test += time.time()
if iter % save_freq == 0:
save_path = "results/dcgan/dcgan.epoch.{}.iter.{}.ckpt". \
format(epoch, iter)
saver.save(sess, save_path)
if iter % print_freq == 0:
time_train = -time.time()
def main():
# Load MNIST
data_path = os.path.join(conf.data_dir, "mnist.pkl.gz")
x_train, t_train, x_valid, t_valid, x_test, t_test = load_mnist_realval(
data_path)
x_train = np.random.binomial(1, x_train,
size=x_train.shape).astype(np.float32)
x_test = np.random.binomial(1, x_test,
size=x_test.shape).astype(np.float32)
x_dim = x_train.shape[1] # x_train = np.vstack([x_train, x_valid])
# Binarize input
y_train = to_categorical(np.array(t_train))
y_test = to_categorical(np.array(t_test))
y_dim = y_train.shape[1]
y_pic = to_categorical(np.arange(10))
# Define model parameters
z_dim = 40
# Build the computation graph
# n_particles = tf.placeholder(tf.int32, shape=[], name="n_particles")
x = tf.placeholder(tf.float32, shape=[None, x_dim], name="x")
y = tf.placeholder(tf.float32, shape=[None, y_dim], name="y")
n = tf.placeholder(tf.int32, shape=[], name="n")
model = build_gen(y, x_dim, z_dim, y_dim, n)
variational = build_q_net(x, y, z_dim, y_dim)
lower_bound = zs.variational.elbo(model, {"x": x}, variational=variational)
cost = tf.reduce_mean(lower_bound.sgvb())
lower_bound = tf.reduce_mean(lower_bound)
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
infer_op = optimizer.minimize(cost)
# is_log_likelihood = tf.reduce_mean(
# zs.is_loglikelihood(model, {"x": x}, proposal=variational, axis=0))
# Random generation
x_gen = tf.reshape(model.observe()["x_mean"], [-1, 28, 28, 1])
epochs = 100
batch_size = 128
iters = x_train.shape[0] // batch_size
# Run the Inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
time_epoch = -time.time()
lbs = []
for t in range(iters):
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
y_batch = y_train[t * batch_size:(t + 1) * batch_size]
_, lb = sess.run([infer_op, lower_bound],
feed_dict={
x: x_batch,
y: y_batch,
n: batch_size
})
lbs.append(lb)
time_epoch += time.time()
print("Epoch {} ({:.1f}s): Lower bound = {}".format(
epoch, time_epoch, np.mean(lbs)))
images = sess.run(x_gen, feed_dict={y: y_pic, n: 10})
name = os.path.join("results", "vae.epoch.{}.png".format(epoch))
save_image_collections(images, name, shape=(1, 10))
def main():
tf.set_random_seed(1234)
np.random.seed(1234)
# Load CIFAR
data_path = os.path.join(conf.data_dir, "cifar10",
"cifar-10-python.tar.gz")
x_train, t_train, x_test, t_test = dataset.load_cifar10(data_path,
normalize=True,
one_hot=True)
# Define model parameters
z_dim = 40
# Build the computation graph
is_training = tf.placeholder(tf.bool, shape=[], name="is_training")
x = tf.placeholder(tf.float32, shape=[None, 32, 32, 3], name="x")
optimizer = tf.train.AdamOptimizer(learning_rate=0.0002, beta1=0.5)
def build_tower_graph(x, id_):
tower_x = x[id_ * tf.shape(x)[0] // FLAGS.num_gpus:(id_ + 1) *
tf.shape(x)[0] // FLAGS.num_gpus]
n = tf.shape(tower_x)[0]
x_gen = generator(n, z_dim, is_training)
x_class_logits = discriminator(tower_x, is_training)
x_gen_class_logits = discriminator(x_gen, is_training)
gen_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=tf.ones_like(x_gen_class_logits),
logits=x_gen_class_logits))
gen_var_list = tf.trainable_variables(scope="gen")
gen_grads = optimizer.compute_gradients(gen_loss,
var_list=gen_var_list)
disc_loss = (tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=tf.ones_like(x_class_logits), logits=x_class_logits)) +
tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=tf.zeros_like(x_gen_class_logits),
logits=x_gen_class_logits))) / 2.
disc_var_list = tf.trainable_variables(scope="disc")
disc_grads = optimizer.compute_gradients(disc_loss,
var_list=disc_var_list)
grads = disc_grads + gen_grads
return grads, gen_loss, disc_loss
tower_losses = []
tower_grads = []
for i in range(FLAGS.num_gpus):
with tf.device("/gpu:%d" % i):
with tf.name_scope("tower_%d" % i):
grads, gen_loss, disc_loss = build_tower_graph(x, i)
tower_losses.append([gen_loss, disc_loss])
tower_grads.append(grads)
gen_loss, disc_loss = multi_gpu.average_losses(tower_losses)
grads = multi_gpu.average_gradients(tower_grads)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
infer_op = optimizer.apply_gradients(grads)
# Generate images
eval_x_gen = generator(100, z_dim, False)
# Define training/evaluation parameters
epochs = 1000
batch_size = 32 * FLAGS.num_gpus
iters = x_train.shape[0] // batch_size
print_freq = 100
save_freq = 100
# Run the inference
with multi_gpu.create_session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
np.random.shuffle(x_train)
gen_losses, disc_losses = [], []
time_train = -time.time()
for t in range(iters):
iter = t + 1
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
_, g_loss, d_loss = sess.run([infer_op, gen_loss, disc_loss],
feed_dict={
x: x_batch,
is_training: True
})
gen_losses.append(g_loss)
disc_losses.append(d_loss)
if iter % print_freq == 0:
print("Epoch={} Iter={} ({:.3f}s/iter): "
"Gen loss = {} Disc loss = {}".format(
epoch, iter,
(time.time() + time_train) / print_freq,
np.mean(gen_losses), np.mean(disc_losses)))
gen_losses = []
disc_losses = []
if iter % save_freq == 0:
images = sess.run(eval_x_gen)
name = "results/dcgan/dcgan.epoch.{}.iter.{}.png".format(
epoch, iter)
save_image_collections(images, name, scale_each=True)
if iter % print_freq == 0:
time_train = -time.time()
def main():
tf.set_random_seed(1234)
np.random.seed(1234)
# Load MNIST
data_path = os.path.join(conf.data_dir, "mnist.pkl.gz")
x_train, t_train, x_valid, t_valid, x_test, t_test = \
dataset.load_mnist_realval(data_path)
x_train = np.vstack([x_train, x_valid])
x_test = np.random.binomial(1, x_test, size=x_test.shape)
x_dim = x_train.shape[1]
# Define model parameters
z_dim = 32
# Build the computation graph
n_particles = tf.placeholder(tf.int32, shape=[], name="n_particles")
x_input = tf.placeholder(tf.float32, shape=[None, x_dim])
x = tf.to_int32(tf.random_uniform(tf.shape(x_input)) <= x_input)
n = tf.shape(x)[0]
def log_joint(observed):
model, _ = vae_conv(observed, n, x_dim, z_dim, n_particles)
log_pz, log_px_z = model.local_log_prob(["z", "x"])
return log_pz + log_px_z
variational = q_net(x, z_dim, n_particles)
qz_samples, log_qz = variational.query("z",
outputs=True,
local_log_prob=True)
lower_bound = zs.variational.elbo(log_joint,
observed={"x": x},
latent={"z": [qz_samples, log_qz]},
axis=0)
cost = tf.reduce_mean(lower_bound.sgvb())
lower_bound = tf.reduce_mean(lower_bound)
optimizer = tf.train.AdamOptimizer(learning_rate=1e-4, beta1=0.5)
infer_op = optimizer.minimize(cost)
# Generate images
n_gen = 100
_, x_logits = vae_conv({}, n_gen, x_dim, z_dim, 1)
x_gen = tf.reshape(tf.sigmoid(x_logits), [-1, 28, 28, 1])
# Define training/evaluation parameters
epochs = 3000
batch_size = 128
iters = x_train.shape[0] // batch_size
save_freq = 10
test_freq = 10
test_batch_size = 400
test_iters = x_test.shape[0] // test_batch_size
result_path = "results/vae_conv"
# Run the inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
time_epoch = -time.time()
np.random.shuffle(x_train)
lbs = []
for t in range(iters):
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
_, lb = sess.run([infer_op, lower_bound],
feed_dict={
x_input: x_batch,
n_particles: 1
})
lbs.append(lb)
time_epoch += time.time()
print("Epoch {} ({:.1f}s): Lower bound = {}".format(
epoch, time_epoch, np.mean(lbs)))
if epoch % test_freq == 0:
time_test = -time.time()
test_lbs = []
for t in range(test_iters):
test_x_batch = x_test[t * test_batch_size:(t + 1) *
test_batch_size]
test_lb = sess.run(lower_bound,
feed_dict={
x: test_x_batch,
n_particles: 1
})
test_lbs.append(test_lb)
time_test += time.time()
print(">>> TEST ({:.1f}s)".format(time_test))
print(">> Test lower bound = {}".format(np.mean(test_lbs)))
if epoch % save_freq == 0:
print("Saving images...")
images = sess.run(x_gen)
name = os.path.join(result_path,
"vae.epoch.{}.png".format(epoch))
save_image_collections(images, name)
def main():
tf.set_random_seed(1234)
np.random.seed(1234)
# Load MINST
data_path = os.path.join(conf.data_dir, 'mnist.pkl.gz')
x_train, t_train, x_valid, t_valid, x_test, t_test = \
dataset.load_mnist_realval(data_path)
n_xl = 28
n_channels = 1
x_train = np.vstack([x_train, x_valid]).astype(np.float32).reshape(
(-1, n_xl, n_xl, n_channels))
# Define model parameters
n_z = 40
# Define training/evaluation parameters
epochs = 1000
batch_size = 64 * FLAGS.num_gpus
gen_size = 100
iters = x_train.shape[0] // batch_size
print_freq = 100
save_freq = 100
# Build the computation graph
is_training = tf.placeholder(tf.bool, shape=[], name='is_training')
x = tf.placeholder(tf.float32,
shape=(None, n_xl, n_xl, n_channels),
name='x')
optimizer = tf.train.RMSPropOptimizer(learning_rate=0.0002, decay=0.5)
def build_tower_graph(x, id_):
tower_x = x[id_ * tf.shape(x)[0] // FLAGS.num_gpus:(id_ + 1) *
tf.shape(x)[0] // FLAGS.num_gpus]
n = tf.shape(tower_x)[0]
gen, x_gen = generator(None, n, n_z, is_training)
x_critic = discriminator(tower_x, is_training)
x_gen_critic = discriminator(x_gen, is_training)
gen_var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='generator')
disc_var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='discriminator')
disc_loss = -tf.reduce_mean(x_critic - x_gen_critic)
gen_loss = -tf.reduce_mean(x_gen_critic)
disc_grads = optimizer.compute_gradients(disc_loss,
var_list=disc_var_list)
gen_grads = optimizer.compute_gradients(gen_loss,
var_list=gen_var_list)
grads = disc_grads + gen_grads
return grads, gen_loss, disc_loss
tower_losses = []
tower_grads = []
for i in range(FLAGS.num_gpus):
with tf.device('/gpu:%d' % i):
with tf.name_scope('tower_%d' % i):
grads, gen_loss, disc_loss = build_tower_graph(x, i)
tower_losses.append([gen_loss, disc_loss])
tower_grads.append(grads)
gen_loss, disc_loss = multi_gpu.average_losses(tower_losses)
w_distance = -disc_loss
grads = multi_gpu.average_gradients(tower_grads)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
infer_op = optimizer.apply_gradients(grads)
# Clip weights of the critic to ensure 1-Lipschitz
disc_var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='discriminator')
with tf.control_dependencies([infer_op]):
clip_op = tf.group(*[
var.assign(tf.clip_by_value(var, -0.01, 0.01))
for var in disc_var_list
])
# Generate images
_, eval_x_gen = generator(None, gen_size, n_z, False)
# Run the inference
with multi_gpu.create_session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
np.random.shuffle(x_train)
w_losses = []
time_train = -time.time()
for t in range(iters):
iter = t + 1
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
_, _, w_loss = sess.run([infer_op, clip_op, w_distance],
feed_dict={
x: x_batch,
is_training: True
})
w_losses.append(w_loss)
if iter % print_freq == 0:
print('Epoch={} Iter={} ({:.3f}s/iter): '
'wasserstein distance = {}'.format(
epoch, iter,
(time.time() + time_train) / print_freq,
np.mean(w_losses)))
w_losses = []
if iter % save_freq == 0:
images = sess.run(eval_x_gen)
name = "results/wgan/wgan.epoch.{}.iter.{}.png".format(
epoch, iter)
save_image_collections(images, name, scale_each=True)
if iter % print_freq == 0:
time_train = -time.time()
def main():
# Load MNIST
data_path = os.path.join(conf.data_dir, 'mnist.pkl.gz')
x_train, t_train, x_valid, t_valid, x_test, t_test = \
dataset.load_mnist_realval(data_path)
x_train = np.random.binomial(1, x_train, size=x_train.shape)
n_x = x_train.shape[1]
# Define model parameters
n_z = 40
@zs.reuse('model')
def vae(observed, n, n_x, n_z):
with zs.BayesianNet(observed=observed) as model:
z_mean = tf.zeros([n, n_z])
z_logstd = tf.zeros([n, n_z])
z = zs.Normal('z', z_mean, logstd=z_logstd, group_event_ndims=1)
lx_z = layers.fully_connected(z, 500)
lx_z = layers.fully_connected(lx_z, 500)
x_logits = layers.fully_connected(lx_z, n_x, activation_fn=None)
x = zs.Bernoulli('x', x_logits, group_event_ndims=1)
return model, x_logits
@zs.reuse('variational')
def q_net(x, n_z):
with zs.BayesianNet() as variational:
lz_x = layers.fully_connected(tf.to_float(x), 500)
lz_x = layers.fully_connected(lz_x, 500)
z_mean = layers.fully_connected(lz_x, n_z, activation_fn=None)
z_logstd = layers.fully_connected(lz_x, n_z, activation_fn=None)
z = zs.Normal('z', z_mean, logstd=z_logstd, group_event_ndims=1)
return variational
x = tf.placeholder(tf.int32, shape=[None, n_x], name='x')
n = tf.shape(x)[0]
def log_joint(observed):
model, _ = vae(observed, n, n_x, n_z)
log_pz, log_px_z = model.local_log_prob(['z', 'x'])
return log_pz + log_px_z
variational = q_net(x, n_z)
qz_samples, log_qz = variational.query('z',
outputs=True,
local_log_prob=True)
lower_bound = tf.reduce_mean(
zs.sgvb(log_joint,
observed={'x': x},
latent={'z': [qz_samples, log_qz]}))
optimizer = tf.train.AdamOptimizer(0.001)
infer = optimizer.minimize(-lower_bound)
# Generate images
n_gen = 100
_, x_logits = vae({}, n_gen, n_x, n_z)
x_gen = tf.reshape(tf.sigmoid(x_logits), [-1, 28, 28, 1])
# Define training parameters
epoches = 500
batch_size = 128
iters = x_train.shape[0] // batch_size
save_freq = 1
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epoches + 1):
np.random.shuffle(x_train)
lbs = []
for t in range(iters):
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
_, lb = sess.run([infer, lower_bound], feed_dict={x: x_batch})
lbs.append(lb)
print('Epoch {}: Lower bound = {}'.format(epoch, np.mean(lbs)))
if epoch % save_freq == 0:
images = sess.run(x_gen)
name = "results/vae/vae.epoch.{}.png".format(epoch)
save_image_collections(images, name)
def main():
# Load MNIST
data_path = os.path.join(conf.data_dir, "mnist.pkl.gz")
x_train, t_train, x_valid, t_valid, x_test, t_test = \
dataset.load_mnist_realval(data_path)
x_train = np.vstack([x_train, x_valid])
x_test = np.random.binomial(1, x_test, size=x_test.shape)
# Define model parameters
x_dim = x_train.shape[1]
z_dim = 40
# Build the computation graph
# how many samples to draw from the distribution, more samples, more accuracy
n_particles = tf.placeholder(tf.int32, shape=[], name="n_particles")
# input data to feed the variational
x_input = tf.placeholder(tf.float32, shape=[None, x_dim], name="x")
x = tf.cast(tf.less(tf.random_uniform(tf.shape(x_input)), x_input),
tf.int32)
# batch size
n = tf.placeholder(tf.int32, shape=[], name="n")
# add random noise to the variance of the q_model so to
# get more various samples when generating new digits
std_noise = tf.placeholder_with_default(0., shape=[], name="std_noise")
# build the model (encoder) and the q_model (variational or decoder)
model = build_gen(x_dim, z_dim, n, n_particles)
q_model = build_q_net(x, z_dim, n_particles, std_noise)
variational = q_model.observe()
# calculate ELBO
lower_bound = zs.variational.elbo(model, {"x": x},
variational=variational,
axis=0)
cost = tf.reduce_mean(lower_bound.sgvb())
lower_bound = tf.reduce_mean(lower_bound)
# calculate marginal log likelihood
is_log_likelihood = tf.reduce_mean(
zs.is_loglikelihood(model, {"x": x}, proposal=variational, axis=0))
# optimizer
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
infer_op = optimizer.minimize(cost)
# define training/evaluation parameters
epochs = 1000
batch_size = 128
iters = x_train.shape[0] // batch_size
test_freq = 100
test_batch_size = 400
test_iters = x_test.shape[0] // test_batch_size
result_path = "results/vae_digits"
checkpoints_path = "checkpoints/vae_digits"
# used to save checkpoints during training
saver = tf.train.Saver(max_to_keep=10)
save_model_freq = 100
# run the inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
# restore the model parameters from the latest checkpoint
ckpt_file = tf.train.latest_checkpoint(checkpoints_path)
begin_epoch = 1
if ckpt_file is not None:
print('Restoring model from {}...'.format(ckpt_file))
begin_epoch = int(ckpt_file.split('.')[-2]) + 1
saver.restore(sess, ckpt_file)
# begin training
for epoch in range(begin_epoch, epochs + 1):
time_epoch = -time.time()
np.random.shuffle(x_train)
lbs = []
for t in range(iters):
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
_, lb = sess.run([infer_op, lower_bound],
feed_dict={
x_input: x_batch,
n_particles: 1,
n: batch_size
})
lbs.append(lb)
time_epoch += time.time()
print("Epoch {} ({:.1f}s): Lower bound = {}".format(
epoch, time_epoch, np.mean(lbs)))
# test marginal log likelihood
if epoch % test_freq == 0:
time_test = -time.time()
test_lbs, test_lls = [], []
for t in range(test_iters):
test_x_batch = x_test[t * test_batch_size:(t + 1) *
test_batch_size]
test_lb = sess.run(lower_bound,
feed_dict={
x: test_x_batch,
n_particles: 1,
n: test_batch_size
})
test_ll = sess.run(is_log_likelihood,
feed_dict={
x: test_x_batch,
n_particles: 1000,
n: test_batch_size
})
test_lbs.append(test_lb)
test_lls.append(test_ll)
time_test += time.time()
print(">>> TEST ({:.1f}s)".format(time_test))
print(">> Test lower bound = {}".format(np.mean(test_lbs)))
print('>> Test log likelihood (IS) = {}'.format(
np.mean(test_lls)))
# save model parameters
if epoch % save_model_freq == 0:
print('Saving model...')
save_path = os.path.join(checkpoints_path,
"vae.epoch.{}.ckpt".format(epoch))
if not os.path.exists(os.path.dirname(save_path)):
os.makedirs(os.path.dirname(save_path))
saver.save(sess, save_path)
print('Done')
# random generation of images from latent distribution
x_gen = tf.reshape(model.observe()["x_mean"], [-1, 28, 28, 1])
images = sess.run(x_gen, feed_dict={n: 100, n_particles: 1})
name = os.path.join(result_path, "random_samples.png")
save_image_collections(images, name)
# the following code generates 100 samples for each number
test_n = [3, 2, 1, 90, 95, 23, 11, 0, 84, 7]
# map each digit to a corresponding sample from the test set so we can generate similar digits
for i in range(len(test_n)):
# get latent distribution from the variational giving as input a fixed sample from the dataset
z = q_model.observe(x=np.expand_dims(x_test[test_n[i]], 0))['z']
# run the computation graph adding noise to computed variance to get different output samples
latent = sess.run(z,
feed_dict={
x_input:
np.expand_dims(x_test[test_n[i]], 0),
n: 1,
n_particles: 100,
std_noise: 0.7
})
# get the image from the model giving as input the latent distribution z
x_gen = tf.reshape(
model.observe(z=latent)["x_mean"], [-1, 28, 28, 1])
images = sess.run(x_gen, feed_dict={})
name = os.path.join(result_path, "{}.png".format(i))
save_image_collections(images, name)
