add a function of visualizing learned manifold for VAE

Additionally
1. correct minor bugs for ACGAN and infoGAN
2. add dimension of z to input-argument
3. save the result images under the folder with name of the simulation settings
4. change batchnorm function as it in tf.contrib.layers

Author: hwalsuklee

ACGAN.py

@@ -9,7 +9,7 @@
 from utils import *
 
 class ACGAN(object):
-    def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result_dir, log_dir):
+    def __init__(self, sess, epoch, batch_size, z_dim, dataset_name, checkpoint_dir, result_dir, log_dir):
         self.sess = sess
         self.dataset_name = dataset_name
         self.checkpoint_dir = checkpoint_dir
@@ -26,7 +26,7 @@ def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result
             self.output_height = 28
             self.output_width = 28
 
-            self.z_dim = 62         # dimension of noise-vector
+            self.z_dim = z_dim         # dimension of noise-vector
             self.y_dim = 10         # dimension of code-vector (label)
             self.c_dim = 1
 
@@ -148,7 +148,7 @@ def build_model(self):
         t_vars = tf.trainable_variables()
         d_vars = [var for var in t_vars if 'd_' in var.name]
         g_vars = [var for var in t_vars if 'g_' in var.name]
-        q_vars = [var for var in t_vars if 'd_' or 'c_' or 'g_' in var.name]
+        q_vars = [var for var in t_vars if ('d_' in var.name) or ('c_' in var.name) or ('g_' in var.name)]
 
         # optimizers
         with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
@@ -246,8 +246,9 @@ def train(self):
                     tot_num_samples = min(self.sample_num, self.batch_size)
                     manifold_h = int(np.floor(np.sqrt(tot_num_samples)))
                     manifold_w = int(np.floor(np.sqrt(tot_num_samples)))
-                    save_images(samples[:manifold_h*manifold_w,:,:,:], [manifold_h, manifold_w],
-                          './'+self.result_dir+'/'+self.model_name+'_train_{:02d}_{:04d}.png'.format(epoch, idx))
+                    save_images(samples[:manifold_h * manifold_w, :, :, :], [manifold_h, manifold_w], './' + check_folder(
+                        self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
+                        epoch, idx))
 
             # After an epoch, start_batch_id is set to zero
             # non-zero value is only for the first epoch after loading pre-trained model
@@ -275,7 +276,7 @@ def visualize_results(self, epoch):
         samples = self.sess.run(self.fake_images, feed_dict={self.inputs:self.test_images, self.z: z_sample, self.y: y_one_hot})
 
         save_images(samples[:image_frame_dim*image_frame_dim,:,:,:], [image_frame_dim, image_frame_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
 
         """ specified condition, random noise """
         n_styles = 10  # must be less than or equal to self.batch_size
@@ -290,7 +291,7 @@ def visualize_results(self, epoch):
 
             samples = self.sess.run(self.fake_images, feed_dict={self.inputs:self.test_images, self.z: z_sample, self.y: y_one_hot})
             save_images(samples[:image_frame_dim*image_frame_dim,:,:,:], [image_frame_dim, image_frame_dim],
-                        self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_class_%d.png' % l)
+                        check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_class_%d.png' % l)
 
             samples = samples[si, :, :, :]
 
@@ -306,13 +307,13 @@ def visualize_results(self, epoch):
                 canvas[s * self.len_discrete_code + c, :, :, :] = all_samples[c * n_styles + s, :, :, :]
 
         save_images(canvas, [n_styles, self.len_discrete_code],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes_style_by_style.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes_style_by_style.png')
 
     @property
     def model_dir(self):
         return "{}_{}_{}_{}".format(
-            self.dataset_name, self.batch_size,
-            self.output_height, self.output_width)
+            self.model_name, self.dataset_name,
+            self.batch_size, self.z_dim)
 
     def save(self, checkpoint_dir, step):
         checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir, self.model_name)

BEGAN.py

@@ -9,7 +9,7 @@
 from utils import *
 
 class BEGAN(object):
-    def __init__(self, sess, batch_size, epoch, dataset_name, checkpoint_dir, result_dir, log_dir):
+    def __init__(self, sess, epoch, batch_size, z_dim, dataset_name, checkpoint_dir, result_dir, log_dir):
         self.sess = sess
         self.dataset_name = dataset_name
         self.checkpoint_dir = checkpoint_dir
@@ -26,7 +26,7 @@ def __init__(self, sess, batch_size, epoch, dataset_name, checkpoint_dir, result
             self.output_height = 28
             self.output_width = 28
 
-            self.z_dim = 62         # dimension of noise-vector
+            self.z_dim = z_dim         # dimension of noise-vector
             self.c_dim = 1
 
             # BEGAN Parameter
@@ -208,7 +208,7 @@ def train(self):
                     manifold_h = int(np.floor(np.sqrt(tot_num_samples)))
                     manifold_w = int(np.floor(np.sqrt(tot_num_samples)))
                     save_images(samples[:manifold_h * manifold_w, :, :, :], [manifold_h, manifold_w],
-                                './' + self.result_dir + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
+                                './' + check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
                                     epoch, idx))
 
             # After an epoch, start_batch_id is set to zero
@@ -235,13 +235,13 @@ def visualize_results(self, epoch):
         samples = self.sess.run(self.fake_images, feed_dict={self.z: z_sample})
 
         save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
 
     @property
     def model_dir(self):
         return "{}_{}_{}_{}".format(
-            self.dataset_name, self.batch_size,
-            self.output_height, self.output_width)
+            self.model_name, self.dataset_name,
+            self.batch_size, self.z_dim)
 
     def save(self, checkpoint_dir, step):
         checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir, self.model_name)

CGAN.py

@@ -9,7 +9,7 @@
 from utils import *
 
 class CGAN(object):
-    def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result_dir, log_dir):
+    def __init__(self, sess, epoch, batch_size, z_dim, dataset_name, checkpoint_dir, result_dir, log_dir):
         self.sess = sess
         self.dataset_name = dataset_name
         self.checkpoint_dir = checkpoint_dir
@@ -26,7 +26,7 @@ def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result
             self.output_height = 28
             self.output_width = 28
 
-            self.z_dim = 62         # dimension of noise-vector
+            self.z_dim = z_dim         # dimension of noise-vector
             self.y_dim = 10         # dimension of condition-vector (label)
             self.c_dim = 1
 
@@ -210,7 +210,7 @@ def train(self):
                     manifold_h = int(np.floor(np.sqrt(tot_num_samples)))
                     manifold_w = int(np.floor(np.sqrt(tot_num_samples)))
                     save_images(samples[:manifold_h * manifold_w, :, :, :], [manifold_h, manifold_w],
-                                './' + self.result_dir + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
+                                './' + check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
                                     epoch, idx))
 
             # After an epoch, start_batch_id is set to zero
@@ -240,7 +240,7 @@ def visualize_results(self, epoch):
         samples = self.sess.run(self.fake_images, feed_dict={self.inputs:self.test_images, self.z: z_sample, self.y: y_one_hot})
 
         save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
 
         """ specified condition, random noise """
         n_styles = 10  # must be less than or equal to self.batch_size
@@ -255,7 +255,7 @@ def visualize_results(self, epoch):
 
             samples = self.sess.run(self.fake_images, feed_dict={self.inputs:self.test_images, self.z: z_sample, self.y: y_one_hot})
             save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],
-                        self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_class_%d.png' % l)
+                        check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_class_%d.png' % l)
 
             samples = samples[si, :, :, :]
 
@@ -271,13 +271,13 @@ def visualize_results(self, epoch):
                 canvas[s * self.y_dim + c, :, :, :] = all_samples[c * n_styles + s, :, :, :]
 
         save_images(canvas, [n_styles, self.y_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes_style_by_style.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes_style_by_style.png')
 
     @property
     def model_dir(self):
         return "{}_{}_{}_{}".format(
-            self.dataset_name, self.batch_size,
-            self.output_height, self.output_width)
+            self.model_name, self.dataset_name,
+            self.batch_size, self.z_dim)
 
     def save(self, checkpoint_dir, step):
         checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir, self.model_name)

CVAE.py

@@ -11,7 +11,7 @@
 import prior_factory as prior
 
 class CVAE(object):
-    def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result_dir, log_dir):
+    def __init__(self, sess, epoch, batch_size, z_dim, dataset_name, checkpoint_dir, result_dir, log_dir):
         self.sess = sess
         self.dataset_name = dataset_name
         self.checkpoint_dir = checkpoint_dir
@@ -28,7 +28,7 @@ def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result
             self.output_height = 28
             self.output_width = 28
 
-            self.z_dim = 62         # dimension of noise-vector
+            self.z_dim = z_dim         # dimension of noise-vector
             self.y_dim = 10         # dimension of condition-vector (label)
             self.c_dim = 1
 
@@ -186,7 +186,7 @@ def train(self):
                 batch_labels = self.data_y[idx * self.batch_size:(idx + 1) * self.batch_size]
                 batch_z = np.random.uniform(-1, 1, [self.batch_size, self.z_dim]).astype(np.float32)
 
-                # update D network
+                # update autoencoder
                 _, summary_str, loss, nll_loss, kl_loss = self.sess.run([self.optim, self.merged_summary_op, self.loss, self.neg_loglikelihood, self.KL_divergence],
                                                feed_dict={self.inputs: batch_images, self.y: batch_labels, self.z: batch_z})
                 self.writer.add_summary(summary_str, counter)
@@ -204,7 +204,7 @@ def train(self):
                     manifold_h = int(np.floor(np.sqrt(tot_num_samples)))
                     manifold_w = int(np.floor(np.sqrt(tot_num_samples)))
                     save_images(samples[:manifold_h * manifold_w, :, :, :], [manifold_h, manifold_w],
-                                './' + self.result_dir + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
+                                './' + check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
                                     epoch, idx))
 
             # After an epoch, start_batch_id is set to zero
@@ -234,7 +234,7 @@ def visualize_results(self, epoch):
         samples = self.sess.run(self.fake_images, feed_dict={self.inputs:self.test_images, self.z: z_sample, self.y: y_one_hot})
 
         save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
 
         """ specified condition, random noise """
         n_styles = 10  # must be less than or equal to self.batch_size
@@ -249,7 +249,7 @@ def visualize_results(self, epoch):
 
             samples = self.sess.run(self.fake_images, feed_dict={self.inputs:self.test_images, self.z: z_sample, self.y: y_one_hot})
             save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],
-                        self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_class_%d.png' % l)
+                        check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_class_%d.png' % l)
 
             samples = samples[si, :, :, :]
 
@@ -265,13 +265,13 @@ def visualize_results(self, epoch):
                 canvas[s * self.y_dim + c, :, :, :] = all_samples[c * n_styles + s, :, :, :]
 
         save_images(canvas, [n_styles, self.y_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes_style_by_style.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes_style_by_style.png')
 
     @property
     def model_dir(self):
         return "{}_{}_{}_{}".format(
-            self.dataset_name, self.batch_size,
-            self.output_height, self.output_width)
+            self.model_name, self.dataset_name,
+            self.batch_size, self.z_dim)
 
     def save(self, checkpoint_dir, step):
         checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir, self.model_name)

DRAGAN.py

@@ -9,7 +9,7 @@
 from utils import *
 
 class DRAGAN(object):
-    def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result_dir, log_dir):
+    def __init__(self, sess, epoch, batch_size, z_dim, dataset_name, checkpoint_dir, result_dir, log_dir):
         self.sess = sess
         self.dataset_name = dataset_name
         self.checkpoint_dir = checkpoint_dir
@@ -26,7 +26,7 @@ def __init__(self, sess, epoch, batch_size, dataset_name, checkpoint_dir, result
             self.output_height = 28
             self.output_width = 28
 
-            self.z_dim = 62         # dimension of noise-vector
+            self.z_dim = z_dim         # dimension of noise-vector
             self.c_dim = 1
 
             # DRAGAN parameter
@@ -211,7 +211,7 @@ def train(self):
                     manifold_h = int(np.floor(np.sqrt(tot_num_samples)))
                     manifold_w = int(np.floor(np.sqrt(tot_num_samples)))
                     save_images(samples[:manifold_h * manifold_w, :, :, :], [manifold_h, manifold_w],
-                                './' + self.result_dir + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
+                                './' + check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
                                     epoch, idx))
 
             # After an epoch, start_batch_id is set to zero
@@ -238,13 +238,13 @@ def visualize_results(self, epoch):
         samples = self.sess.run(self.fake_images, feed_dict={self.z: z_sample})
 
         save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
 
     @property
     def model_dir(self):
         return "{}_{}_{}_{}".format(
-            self.dataset_name, self.batch_size,
-            self.output_height, self.output_width)
+            self.model_name, self.dataset_name,
+            self.batch_size, self.z_dim)
 
     def save(self, checkpoint_dir, step):
         checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir, self.model_name)

EBGAN.py

@@ -9,7 +9,7 @@
 from utils import *
 
 class EBGAN(object):
-    def __init__(self, sess, batch_size, epoch, dataset_name, checkpoint_dir, result_dir, log_dir):
+    def __init__(self, sess, epoch, batch_size, z_dim, dataset_name, checkpoint_dir, result_dir, log_dir):
         self.sess = sess
         self.dataset_name = dataset_name
         self.checkpoint_dir = checkpoint_dir
@@ -26,7 +26,7 @@ def __init__(self, sess, batch_size, epoch, dataset_name, checkpoint_dir, result
             self.output_height = 28
             self.output_width = 28
 
-            self.z_dim = 62         # dimension of noise-vector
+            self.z_dim = z_dim         # dimension of noise-vector
             self.c_dim = 1
 
             # EBGAN Parameter
@@ -208,7 +208,7 @@ def train(self):
                     manifold_h = int(np.floor(np.sqrt(tot_num_samples)))
                     manifold_w = int(np.floor(np.sqrt(tot_num_samples)))
                     save_images(samples[:manifold_h * manifold_w, :, :, :], [manifold_h, manifold_w],
-                                './' + self.result_dir + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
+                                './' + check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
                                     epoch, idx))
 
             # After an epoch, start_batch_id is set to zero
@@ -235,13 +235,13 @@ def visualize_results(self, epoch):
         samples = self.sess.run(self.fake_images, feed_dict={self.z: z_sample})
 
         save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],
-                    self.result_dir + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
+                    check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_epoch%03d' % epoch + '_test_all_classes.png')
 
     @property
     def model_dir(self):
         return "{}_{}_{}_{}".format(
-            self.dataset_name, self.batch_size,
-            self.output_height, self.output_width)
+            self.model_name, self.dataset_name,
+            self.batch_size, self.z_dim)
 
     def save(self, checkpoint_dir, step):
         checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir, self.model_name)