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)