upgrade to tf 1.1.4; slim <- keras; tf.distributions <- tf.probability; tf <- tf.compat.v1

Author: Jaan Altosaar

environment.yml

@@ -0,0 +1,74 @@
+name: dev
+channels:
+  - defaults
+dependencies:
+  - blas=1.0=mkl
+  - ca-certificates=2019.5.15=1
+  - certifi=2019.6.16=py37_1
+  - freetype=2.9.1=hb4e5f40_0
+  - imageio=2.5.0=py37_0
+  - intel-openmp=2019.4=233
+  - jpeg=9b=he5867d9_2
+  - libcxx=4.0.1=hcfea43d_1
+  - libcxxabi=4.0.1=hcfea43d_1
+  - libedit=3.1.20181209=hb402a30_0
+  - libffi=3.2.1=h475c297_4
+  - libgfortran=3.0.1=h93005f0_2
+  - libpng=1.6.37=ha441bb4_0
+  - libtiff=4.0.10=hcb84e12_2
+  - mkl=2019.4=233
+  - mkl-service=2.3.0=py37hfbe908c_0
+  - mkl_fft=1.0.14=py37h5e564d8_0
+  - mkl_random=1.0.2=py37h27c97d8_0
+  - ncurses=6.1=h0a44026_1
+  - numpy=1.16.5=py37hacdab7b_0
+  - numpy-base=1.16.5=py37h6575580_0
+  - olefile=0.46=py37_0
+  - openssl=1.1.1d=h1de35cc_1
+  - pillow=6.1.0=py37hb68e598_0
+  - python=3.7.4=h359304d_1
+  - readline=7.0=h1de35cc_5
+  - setuptools=41.0.1=py37_0
+  - six=1.12.0=py37_0
+  - sqlite=3.29.0=ha441bb4_0
+  - tk=8.6.8=ha441bb4_0
+  - wheel=0.33.4=py37_0
+  - xz=5.2.4=h1de35cc_4
+  - zlib=1.2.11=h1de35cc_3
+  - zstd=1.3.7=h5bba6e5_0
+  - pip:
+    - absl-py==0.8.0
+    - astor==0.8.0
+    - attrs==19.1.0
+    - chardet==3.0.4
+    - cloudpickle==1.2.2
+    - decorator==4.4.0
+    - dill==0.3.0
+    - future==0.17.1
+    - gast==0.3.2
+    - google-pasta==0.1.7
+    - googleapis-common-protos==1.6.0
+    - grpcio==1.23.0
+    - h5py==2.10.0
+    - idna==2.8
+    - keras-applications==1.0.8
+    - keras-preprocessing==1.1.0
+    - markdown==3.1.1
+    - pip==19.2.3
+    - promise==2.2.1
+    - protobuf==3.9.1
+    - psutil==5.6.3
+    - requests==2.22.0
+    - tensorboard==1.14.0
+    - tensorflow==1.14.0
+    - tensorflow-datasets==1.2.0
+    - tensorflow-estimator==1.14.0
+    - tensorflow-metadata==0.14.0
+    - tensorflow-probability==0.7.0
+    - termcolor==1.1.0
+    - tqdm==4.36.0
+    - urllib3==1.25.3
+    - werkzeug==0.15.6
+    - wrapt==1.11.2
+prefix: /usr/local/anaconda3/envs/dev
+

train_variational_autoencoder_tensorflow.py

@@ -1,33 +1,20 @@
 import itertools
-import matplotlib as mpl
 import numpy as np
 import os
 import tensorflow as tf
+import tensorflow.keras as tfk
 import tensorflow.contrib.slim as slim
 import time
-import seaborn as sns
-
-from matplotlib import pyplot as plt
+import tensorflow_datasets as tfds
+import tensorflow_probability as tfp
 from imageio import imwrite
 from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets
-
-sns.set_style('whitegrid')
-
-distributions = tf.distributions
+tfkl = tfk.layers
+tfc = tf.compat.v1
 
 flags = tf.app.flags
 flags.DEFINE_string('data_dir', '/tmp/dat/', 'Directory for data')
 flags.DEFINE_string('logdir', '/tmp/log/', 'Directory for logs')
-
-# For making plots:
-# flags.DEFINE_integer('latent_dim', 2, 'Latent dimensionality of model')
-# flags.DEFINE_integer('batch_size', 64, 'Minibatch size')
-# flags.DEFINE_integer('n_samples', 10, 'Number of samples to save')
-# flags.DEFINE_integer('print_every', 10, 'Print every n iterations')
-# flags.DEFINE_integer('hidden_size', 200, 'Hidden size for neural networks')
-# flags.DEFINE_integer('n_iterations', 1000, 'number of iterations')
-
-# For bigger model:
 flags.DEFINE_integer('latent_dim', 100, 'Latent dimensionality of model')
 flags.DEFINE_integer('batch_size', 64, 'Minibatch size')
 flags.DEFINE_integer('n_samples', 1, 'Number of samples to save')
@@ -50,12 +37,13 @@ def inference_network(x, latent_dim, hidden_size):
     mu: Mean parameters for the variational family Normal
     sigma: Standard deviation parameters for the variational family Normal
   """
-  with slim.arg_scope([slim.fully_connected], activation_fn=tf.nn.relu):
-    net = slim.flatten(x)
-    net = slim.fully_connected(net, hidden_size)
-    net = slim.fully_connected(net, hidden_size)
-    gaussian_params = slim.fully_connected(
-        net, latent_dim * 2, activation_fn=None)
+  inference_net = tfk.Sequential([
+    tfkl.Flatten(),
+    tfkl.Dense(hidden_size, activation=tf.nn.relu),
+    tfkl.Dense(hidden_size, activation=tf.nn.relu),
+    tfkl.Dense(latent_dim * 2, activation=None)
+    ])
+  gaussian_params = inference_net(x)
   # The mean parameter is unconstrained
   mu = gaussian_params[:, :latent_dim]
   # The standard deviation must be positive. Parametrize with a softplus
@@ -73,174 +61,111 @@ def generative_network(z, hidden_size):
   Returns:
     bernoulli_logits: logits for the Bernoulli likelihood of the data
   """
-  with slim.arg_scope([slim.fully_connected], activation_fn=tf.nn.relu):
-    net = slim.fully_connected(z, hidden_size)
-    net = slim.fully_connected(net, hidden_size)
-    bernoulli_logits = slim.fully_connected(net, 784, activation_fn=None)
-    bernoulli_logits = tf.reshape(bernoulli_logits, [-1, 28, 28, 1])
-  return bernoulli_logits
+  generative_net = tfk.Sequential([
+    tfkl.Dense(hidden_size, activation=tf.nn.relu),
+    tfkl.Dense(hidden_size, activation=tf.nn.relu),
+    tfkl.Dense(28 * 28, activation=None)
+    ])
+  bernoulli_logits = generative_net(z)
+  return tf.reshape(bernoulli_logits, [-1, 28, 28, 1])
 
 
 def train():
   # Train a Variational Autoencoder on MNIST
 
   # Input placeholders
   with tf.name_scope('data'):
-    x = tf.placeholder(tf.float32, [None, 28, 28, 1])
-    tf.summary.image('data', x)
+    x = tfc.placeholder(tf.float32, [None, 28, 28, 1])
+    tfc.summary.image('data', x)
 
-  with tf.variable_scope('variational'):
+  with tfc.variable_scope('variational'):
     q_mu, q_sigma = inference_network(x=x,
                                       latent_dim=FLAGS.latent_dim,
                                       hidden_size=FLAGS.hidden_size)
     # The variational distribution is a Normal with mean and standard
     # deviation given by the inference network
-    q_z = distributions.Normal(loc=q_mu, scale=q_sigma)
-    assert q_z.reparameterization_type == distributions.FULLY_REPARAMETERIZED
+    q_z = tfp.distributions.Normal(loc=q_mu, scale=q_sigma)
+    assert q_z.reparameterization_type == tfp.distributions.FULLY_REPARAMETERIZED
 
-  with tf.variable_scope('model'):
+  with tfc.variable_scope('model'):
     # The likelihood is Bernoulli-distributed with logits given by the
     # generative network
     p_x_given_z_logits = generative_network(z=q_z.sample(),
                                             hidden_size=FLAGS.hidden_size)
-    p_x_given_z = distributions.Bernoulli(logits=p_x_given_z_logits)
+    p_x_given_z = tfp.distributions.Bernoulli(logits=p_x_given_z_logits)
     posterior_predictive_samples = p_x_given_z.sample()
-    tf.summary.image('posterior_predictive',
+    tfc.summary.image('posterior_predictive',
                      tf.cast(posterior_predictive_samples, tf.float32))
 
   # Take samples from the prior
-  with tf.variable_scope('model', reuse=True):
-    p_z = distributions.Normal(loc=np.zeros(FLAGS.latent_dim, dtype=np.float32),
+  with tfc.variable_scope('model', reuse=True):
+    p_z = tfp.distributions.Normal(loc=np.zeros(FLAGS.latent_dim, dtype=np.float32),
                                scale=np.ones(FLAGS.latent_dim, dtype=np.float32))
     p_z_sample = p_z.sample(FLAGS.n_samples)
     p_x_given_z_logits = generative_network(z=p_z_sample,
                                             hidden_size=FLAGS.hidden_size)
-    prior_predictive = distributions.Bernoulli(logits=p_x_given_z_logits)
+    prior_predictive = tfp.distributions.Bernoulli(logits=p_x_given_z_logits)
     prior_predictive_samples = prior_predictive.sample()
-    tf.summary.image('prior_predictive',
+    tfc.summary.image('prior_predictive',
                      tf.cast(prior_predictive_samples, tf.float32))
 
   # Take samples from the prior with a placeholder
-  with tf.variable_scope('model', reuse=True):
+  with tfc.variable_scope('model', reuse=True):
     z_input = tf.placeholder(tf.float32, [None, FLAGS.latent_dim])
     p_x_given_z_logits = generative_network(z=z_input,
                                             hidden_size=FLAGS.hidden_size)
-    prior_predictive_inp = distributions.Bernoulli(logits=p_x_given_z_logits)
+    prior_predictive_inp = tfp.distributions.Bernoulli(logits=p_x_given_z_logits)
     prior_predictive_inp_sample = prior_predictive_inp.sample()
 
   # Build the evidence lower bound (ELBO) or the negative loss
-  kl = tf.reduce_sum(distributions.kl_divergence(q_z, p_z), 1)
+  kl = tf.reduce_sum(tfp.distributions.kl_divergence(q_z, p_z), 1)
   expected_log_likelihood = tf.reduce_sum(p_x_given_z.log_prob(x),
                                           [1, 2, 3])
 
   elbo = tf.reduce_sum(expected_log_likelihood - kl, 0)
-
-  optimizer = tf.train.RMSPropOptimizer(learning_rate=0.001)
-
+  optimizer = tfc.train.RMSPropOptimizer(learning_rate=0.001)
   train_op = optimizer.minimize(-elbo)
 
   # Merge all the summaries
-  summary_op = tf.summary.merge_all()
+  summary_op = tfc.summary.merge_all()
 
-  init_op = tf.global_variables_initializer()
+  init_op = tfc.global_variables_initializer()
 
   # Run training
-  sess = tf.InteractiveSession()
+  sess = tfc.InteractiveSession()
   sess.run(init_op)
 
-  mnist = read_data_sets(FLAGS.data_dir, one_hot=True)
+  mnist_data = tfds.load(name='binarized_mnist', split='train', shuffle_files=False)
+  dataset = mnist_data.repeat().shuffle(buffer_size=1024).batch(FLAGS.batch_size)
 
   print('Saving TensorBoard summaries and images to: %s' % FLAGS.logdir)
-  train_writer = tf.summary.FileWriter(FLAGS.logdir, sess.graph)
-
-  # Get fixed MNIST digits for plotting posterior means during training
-  np_x_fixed, np_y = mnist.test.next_batch(5000)
-  np_x_fixed = np_x_fixed.reshape(5000, 28, 28, 1)
-  np_x_fixed = (np_x_fixed > 0.5).astype(np.float32)
+  train_writer = tfc.summary.FileWriter(FLAGS.logdir, sess.graph)
 
   t0 = time.time()
-  for i in range(FLAGS.n_iterations):
-    # Re-binarize the data at every batch; this improves results
-    np_x, _ = mnist.train.next_batch(FLAGS.batch_size)
-    np_x = np_x.reshape(FLAGS.batch_size, 28, 28, 1)
-    np_x = (np_x > 0.5).astype(np.float32)
+  for i, batch in enumerate(tfds.as_numpy(dataset)):
+    np_x = batch['image']
     sess.run(train_op, {x: np_x})
-
-    # Print progress and save samples every so often
     if i % FLAGS.print_every == 0:
       np_elbo, summary_str = sess.run([elbo, summary_op], {x: np_x})
       train_writer.add_summary(summary_str, i)
       print('Iteration: {0:d} ELBO: {1:.3f} s/iter: {2:.3e}'.format(
           i,
           np_elbo / FLAGS.batch_size,
           (time.time() - t0) / FLAGS.print_every))
-      t0 = time.time()
-
       # Save samples
       np_posterior_samples, np_prior_samples = sess.run(
           [posterior_predictive_samples, prior_predictive_samples], {x: np_x})
       for k in range(FLAGS.n_samples):
         f_name = os.path.join(
             FLAGS.logdir, 'iter_%d_posterior_predictive_%d_data.jpg' % (i, k))
-        imwrite(f_name, np_x[k, :, :, 0])
+        imwrite(f_name, np_x[k, :, :, 0].astype(np.uint8))
         f_name = os.path.join(
             FLAGS.logdir, 'iter_%d_posterior_predictive_%d_sample.jpg' % (i, k))
-        imwrite(f_name, np_posterior_samples[k, :, :, 0])
+        imwrite(f_name, np_posterior_samples[k, :, :, 0].astype(np.uint8))
         f_name = os.path.join(
             FLAGS.logdir, 'iter_%d_prior_predictive_%d.jpg' % (i, k))
-        imwrite(f_name, np_prior_samples[k, :, :, 0])
-
-      # Plot the posterior predictive space
-      if FLAGS.latent_dim == 2:
-        np_q_mu = sess.run(q_mu, {x: np_x_fixed})
-        cmap = mpl.colors.ListedColormap(sns.color_palette("husl"))
-        f, ax = plt.subplots(1, figsize=(6 * 1.1618, 6))
-        im = ax.scatter(np_q_mu[:, 0], np_q_mu[:, 1], c=np.argmax(np_y, 1), cmap=cmap,
-                        alpha=0.7)
-        ax.set_xlabel('First dimension of sampled latent variable $z_1$')
-        ax.set_ylabel('Second dimension of sampled latent variable mean $z_2$')
-        ax.set_xlim([-10., 10.])
-        ax.set_ylim([-10., 10.])
-        f.colorbar(im, ax=ax, label='Digit class')
-        plt.tight_layout()
-        plt.savefig(os.path.join(FLAGS.logdir,
-                                 'posterior_predictive_map_frame_%d.png' % i))
-        plt.close()
-
-        nx = ny = 20
-        x_values = np.linspace(-3, 3, nx)
-        y_values = np.linspace(-3, 3, ny)
-        canvas = np.empty((28 * ny, 28 * nx))
-        for ii, yi in enumerate(x_values):
-          for j, xi in enumerate(y_values):
-            np_z = np.array([[xi, yi]])
-            x_mean = sess.run(prior_predictive_inp_sample, {z_input: np_z})
-            canvas[(nx - ii - 1) * 28:(nx - ii) * 28, j *
-                   28:(j + 1) * 28] = x_mean[0].reshape(28, 28)
-        imwrite(os.path.join(FLAGS.logdir,
-                            'prior_predictive_map_frame_%d.png' % i), canvas)
-        # plt.figure(figsize=(8, 10))
-        # Xi, Yi = np.meshgrid(x_values, y_values)
-        # plt.imshow(canvas, origin="upper")
-        # plt.tight_layout()
-        # plt.savefig()
-
-  # Make the gifs
-  if FLAGS.latent_dim == 2:
-    os.system(
-        'convert -delay 15 -loop 0 {0}/posterior_predictive_map_frame*png {0}/posterior_predictive.gif'
-        .format(FLAGS.logdir))
-    os.system(
-        'convert -delay 15 -loop 0 {0}/prior_predictive_map_frame*png {0}/prior_predictive.gif'
-        .format(FLAGS.logdir))
-
-
-def main(_):
-  if tf.gfile.Exists(FLAGS.logdir):
-    tf.gfile.DeleteRecursively(FLAGS.logdir)
-  tf.gfile.MakeDirs(FLAGS.logdir)
-  train()
-
+        imwrite(f_name, np_prior_samples[k, :, :, 0].astype(np.uint8))
+      t0 = time.time()
 
 if __name__ == '__main__':
-  tf.app.run()
+  train()