Merge pull request #302 from memo/master

A number of extra options

Author: carpedm20

main.py

@@ -1,9 +1,10 @@
 import os
 import scipy.misc
 import numpy as np
+import json
 
 from model import DCGAN
-from utils import pp, visualize, to_json, show_all_variables
+from utils import pp, visualize, to_json, show_all_variables, expand_path, timestamp
 
 import tensorflow as tf
 
@@ -19,27 +20,56 @@
 flags.DEFINE_integer("output_width", None, "The size of the output images to produce. If None, same value as output_height [None]")
 flags.DEFINE_string("dataset", "celebA", "The name of dataset [celebA, mnist, lsun]")
 flags.DEFINE_string("input_fname_pattern", "*.jpg", "Glob pattern of filename of input images [*]")
-flags.DEFINE_string("checkpoint_dir", "checkpoint", "Directory name to save the checkpoints [checkpoint]")
-flags.DEFINE_string("data_dir", "./data", "Root directory of dataset [data]")
-flags.DEFINE_string("sample_dir", "samples", "Directory name to save the image samples [samples]")
+flags.DEFINE_string("data_dir", "./data", "path to datasets [e.g. $HOME/data]")
+flags.DEFINE_string("out_dir", "./out", "Root directory for outputs [e.g. $HOME/out]")
+flags.DEFINE_string("out_name", "", "Folder (under out_root_dir) for all outputs. Generated automatically if left blank []")
+flags.DEFINE_string("checkpoint_dir", "checkpoint", "Folder (under out_root_dir/out_name) to save checkpoints [checkpoint]")
+flags.DEFINE_string("sample_dir", "samples", "Folder (under out_root_dir/out_name) to save samples [samples]")
 flags.DEFINE_boolean("train", False, "True for training, False for testing [False]")
 flags.DEFINE_boolean("crop", False, "True for training, False for testing [False]")
 flags.DEFINE_boolean("visualize", False, "True for visualizing, False for nothing [False]")
-flags.DEFINE_integer("generate_test_images", 100, "Number of images to generate during test. [100]")
+flags.DEFINE_boolean("export", False, "True for exporting with new batch size")
+flags.DEFINE_boolean("freeze", False, "True for exporting with new batch size")
+flags.DEFINE_integer("max_to_keep", 1, "maximum number of checkpoints to keep")
+flags.DEFINE_integer("sample_freq", 200, "sample every this many iterations")
+flags.DEFINE_integer("ckpt_freq", 200, "save checkpoint every this many iterations")
+flags.DEFINE_integer("z_dim", 100, "dimensions of z")
+flags.DEFINE_string("z_dist", "uniform_signed", "'normal01' or 'uniform_unsigned' or uniform_signed")
+flags.DEFINE_boolean("G_img_sum", False, "Save generator image summaries in log")
+#flags.DEFINE_integer("generate_test_images", 100, "Number of images to generate during test. [100]")
 FLAGS = flags.FLAGS
 
 def main(_):
   pp.pprint(flags.FLAGS.__flags)
+  
+  # expand user name and environment variables
+  FLAGS.data_dir = expand_path(FLAGS.data_dir)
+  FLAGS.out_dir = expand_path(FLAGS.out_dir)
+  FLAGS.out_name = expand_path(FLAGS.out_name)
+  FLAGS.checkpoint_dir = expand_path(FLAGS.checkpoint_dir)
+  FLAGS.sample_dir = expand_path(FLAGS.sample_dir)
 
-  if FLAGS.input_width is None:
-    FLAGS.input_width = FLAGS.input_height
-  if FLAGS.output_width is None:
-    FLAGS.output_width = FLAGS.output_height
+  if FLAGS.output_height is None: FLAGS.output_height = FLAGS.input_height
+  if FLAGS.input_width is None: FLAGS.input_width = FLAGS.input_height
+  if FLAGS.output_width is None: FLAGS.output_width = FLAGS.output_height
 
-  if not os.path.exists(FLAGS.checkpoint_dir):
-    os.makedirs(FLAGS.checkpoint_dir)
-  if not os.path.exists(FLAGS.sample_dir):
-    os.makedirs(FLAGS.sample_dir)
+  # output folders
+  if FLAGS.out_name == "":
+      FLAGS.out_name = '{} - {} - {}'.format(timestamp(), FLAGS.data_dir.split('/')[-1], FLAGS.dataset) # penultimate folder of path
+      if FLAGS.train:
+        FLAGS.out_name += ' - x{}.z{}.{}.y{}.b{}'.format(FLAGS.input_width, FLAGS.z_dim, FLAGS.z_dist, FLAGS.output_width, FLAGS.batch_size)
+
+  FLAGS.out_dir = os.path.join(FLAGS.out_dir, FLAGS.out_name)
+  FLAGS.checkpoint_dir = os.path.join(FLAGS.out_dir, FLAGS.checkpoint_dir)
+  FLAGS.sample_dir = os.path.join(FLAGS.out_dir, FLAGS.sample_dir)
+
+  if not os.path.exists(FLAGS.checkpoint_dir): os.makedirs(FLAGS.checkpoint_dir)
+  if not os.path.exists(FLAGS.sample_dir): os.makedirs(FLAGS.sample_dir)
+
+  with open(os.path.join(FLAGS.out_dir, 'FLAGS.json'), 'w') as f:
+    flags_dict = {k:FLAGS[k].value for k in FLAGS}
+    json.dump(flags_dict, f, indent=4, sort_keys=True, ensure_ascii=False)
+  
 
   #gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
   run_config = tf.ConfigProto()
@@ -56,13 +86,15 @@ def main(_):
           batch_size=FLAGS.batch_size,
           sample_num=FLAGS.batch_size,
           y_dim=10,
-          z_dim=FLAGS.generate_test_images,
+          z_dim=FLAGS.z_dim,
           dataset_name=FLAGS.dataset,
           input_fname_pattern=FLAGS.input_fname_pattern,
           crop=FLAGS.crop,
           checkpoint_dir=FLAGS.checkpoint_dir,
           sample_dir=FLAGS.sample_dir,
-          data_dir=FLAGS.data_dir)
+          data_dir=FLAGS.data_dir,
+          out_dir=FLAGS.out_dir,
+          max_to_keep=FLAGS.max_to_keep)
     else:
       dcgan = DCGAN(
           sess,
@@ -72,22 +104,25 @@ def main(_):
           output_height=FLAGS.output_height,
           batch_size=FLAGS.batch_size,
           sample_num=FLAGS.batch_size,
-          z_dim=FLAGS.generate_test_images,
+          z_dim=FLAGS.z_dim,
           dataset_name=FLAGS.dataset,
           input_fname_pattern=FLAGS.input_fname_pattern,
           crop=FLAGS.crop,
           checkpoint_dir=FLAGS.checkpoint_dir,
           sample_dir=FLAGS.sample_dir,
-          data_dir=FLAGS.data_dir)
+          data_dir=FLAGS.data_dir,
+          out_dir=FLAGS.out_dir,
+          max_to_keep=FLAGS.max_to_keep)
 
     show_all_variables()
 
     if FLAGS.train:
       dcgan.train(FLAGS)
     else:
-      if not dcgan.load(FLAGS.checkpoint_dir)[0]:
-        raise Exception("[!] Train a model first, then run test mode")
-      
+      load_success, load_counter = dcgan.load(FLAGS.checkpoint_dir)
+      if not load_success:
+        raise Exception("Checkpoint not found in " + FLAGS.checkpoint_dir)
+
 
     # to_json("./web/js/layers.js", [dcgan.h0_w, dcgan.h0_b, dcgan.g_bn0],
     #                 [dcgan.h1_w, dcgan.h1_b, dcgan.g_bn1],
@@ -96,8 +131,17 @@ def main(_):
     #                 [dcgan.h4_w, dcgan.h4_b, None])
 
     # Below is codes for visualization
-    OPTION = 1
-    visualize(sess, dcgan, FLAGS, OPTION)
+      if FLAGS.export:
+        export_dir = os.path.join(FLAGS.checkpoint_dir, 'export_b'+str(FLAGS.batch_size))
+        dcgan.save(export_dir, load_counter, ckpt=True, frozen=False)
+
+      if FLAGS.freeze:
+        export_dir = os.path.join(FLAGS.checkpoint_dir, 'frozen_b'+str(FLAGS.batch_size))
+        dcgan.save(export_dir, load_counter, ckpt=False, frozen=True)
+
+      if FLAGS.visualize:
+        OPTION = 1
+        visualize(sess, dcgan, FLAGS, OPTION, FLAGS.sample_dir)
 
 if __name__ == '__main__':
   tf.app.run()

model.py

@@ -1,4 +1,5 @@
 from __future__ import division
+from __future__ import print_function
 import os
 import time
 import math
@@ -13,12 +14,19 @@
 def conv_out_size_same(size, stride):
   return int(math.ceil(float(size) / float(stride)))
 
+def gen_random(mode, size):
+    if mode=='normal01': return np.random.normal(0,1,size=size)
+    if mode=='uniform_signed': return np.random.uniform(-1,1,size=size)
+    if mode=='uniform_unsigned': return np.random.uniform(0,1,size=size)
+
+
 class DCGAN(object):
   def __init__(self, sess, input_height=108, input_width=108, crop=True,
          batch_size=64, sample_num = 64, output_height=64, output_width=64,
          y_dim=None, z_dim=100, gf_dim=64, df_dim=64,
          gfc_dim=1024, dfc_dim=1024, c_dim=3, dataset_name='default',
-         input_fname_pattern='*.jpg', checkpoint_dir=None, sample_dir=None, data_dir='./data'):
+         max_to_keep=1,
+         input_fname_pattern='*.jpg', checkpoint_dir='ckpts', sample_dir='samples', out_dir='./out', data_dir='./data'):
     """
 
     Args:
@@ -70,6 +78,8 @@ def __init__(self, sess, input_height=108, input_width=108, crop=True,
     self.input_fname_pattern = input_fname_pattern
     self.checkpoint_dir = checkpoint_dir
     self.data_dir = data_dir
+    self.out_dir = out_dir
+    self.max_to_keep = max_to_keep
 
     if self.dataset_name == 'mnist':
       self.data_X, self.data_y = self.load_mnist()
@@ -148,7 +158,7 @@ def sigmoid_cross_entropy_with_logits(x, y):
     self.d_vars = [var for var in t_vars if 'd_' in var.name]
     self.g_vars = [var for var in t_vars if 'g_' in var.name]
 
-    self.saver = tf.train.Saver()
+    self.saver = tf.train.Saver(max_to_keep=self.max_to_keep)
 
   def train(self, config):
     d_optim = tf.train.AdamOptimizer(config.learning_rate, beta1=config.beta1) \
@@ -160,13 +170,15 @@ def train(self, config):
     except:
       tf.initialize_all_variables().run()
 
-    self.g_sum = merge_summary([self.z_sum, self.d__sum,
-      self.G_sum, self.d_loss_fake_sum, self.g_loss_sum])
+    if config.G_img_sum:
+      self.g_sum = merge_summary([self.z_sum, self.d__sum, self.G_sum, self.d_loss_fake_sum, self.g_loss_sum])
+    else:
+      self.g_sum = merge_summary([self.z_sum, self.d__sum, self.d_loss_fake_sum, self.g_loss_sum])
     self.d_sum = merge_summary(
         [self.z_sum, self.d_sum, self.d_loss_real_sum, self.d_loss_sum])
-    self.writer = SummaryWriter("./logs", self.sess.graph)
+    self.writer = SummaryWriter(os.path.join(self.out_dir, "logs"), self.sess.graph)
 
-    sample_z = np.random.uniform(-1, 1, size=(self.sample_num , self.z_dim))
+    sample_z = gen_random(config.z_dist, size=(self.sample_num , self.z_dim))
 
     if config.dataset == 'mnist':
       sample_inputs = self.data_X[0:self.sample_num]
@@ -223,7 +235,7 @@ def train(self, config):
           else:
             batch_images = np.array(batch).astype(np.float32)
 
-        batch_z = np.random.uniform(-1, 1, [config.batch_size, self.z_dim]) \
+        batch_z = gen_random(config.z_dist, size=[config.batch_size, self.z_dim]) \
               .astype(np.float32)
 
         if config.dataset == 'mnist':
@@ -281,12 +293,11 @@ def train(self, config):
           errD_real = self.d_loss_real.eval({ self.inputs: batch_images })
           errG = self.g_loss.eval({self.z: batch_z})
 
-        counter += 1
-        print("Epoch: [%2d/%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" \
-          % (epoch, config.epoch, idx, batch_idxs,
+        print("[%8d Epoch:[%2d/%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" \
+          % (counter, epoch, config.epoch, idx, batch_idxs,
             time.time() - start_time, errD_fake+errD_real, errG))
 
-        if np.mod(counter, 100) == 1:
+        if np.mod(counter, config.sample_freq) == 0:
           if config.dataset == 'mnist':
             samples, d_loss, g_loss = self.sess.run(
               [self.sampler, self.d_loss, self.g_loss],
@@ -297,7 +308,7 @@ def train(self, config):
               }
             )
             save_images(samples, image_manifold_size(samples.shape[0]),
-                  './{}/train_{:02d}_{:04d}.png'.format(config.sample_dir, epoch, idx))
+                  './{}/train_{:08d}.png'.format(config.sample_dir, counter))
             print("[Sample] d_loss: %.8f, g_loss: %.8f" % (d_loss, g_loss)) 
           else:
             try:
@@ -309,14 +320,16 @@ def train(self, config):
                 },
               )
               save_images(samples, image_manifold_size(samples.shape[0]),
-                    './{}/train_{:02d}_{:04d}.png'.format(config.sample_dir, epoch, idx))
+                    './{}/train_{:08d}.png'.format(config.sample_dir, counter))
               print("[Sample] d_loss: %.8f, g_loss: %.8f" % (d_loss, g_loss)) 
             except:
               print("one pic error!...")
 
-        if np.mod(counter, 500) == 2:
+        if np.mod(counter, config.ckpt_freq) == 0:
           self.save(config.checkpoint_dir, counter)
-
+        
+        counter += 1
+        
   def discriminator(self, image, y=None, reuse=False):
     with tf.variable_scope("discriminator") as scope:
       if reuse:
@@ -501,28 +514,39 @@ def model_dir(self):
     return "{}_{}_{}_{}".format(
         self.dataset_name, self.batch_size,
         self.output_height, self.output_width)
-      
-  def save(self, checkpoint_dir, step):
-    model_name = "DCGAN.model"
-    checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)
 
+  def save(self, checkpoint_dir, step, filename='model', ckpt=True, frozen=False):
+    # model_name = "DCGAN.model"
+    # checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)
+
+    filename += '.b' + str(self.batch_size)
     if not os.path.exists(checkpoint_dir):
       os.makedirs(checkpoint_dir)
 
-    self.saver.save(self.sess,
-            os.path.join(checkpoint_dir, model_name),
-            global_step=step)
+    if ckpt:
+      self.saver.save(self.sess,
+              os.path.join(checkpoint_dir, filename),
+              global_step=step)
+
+    if frozen:
+      tf.train.write_graph(
+              tf.graph_util.convert_variables_to_constants(self.sess, self.sess.graph_def, ["generator_1/Tanh"]),
+              checkpoint_dir,
+              '{}-{:06d}_frz.pb'.format(filename, step),
+              as_text=False)
 
   def load(self, checkpoint_dir):
-    import re
-    print(" [*] Reading checkpoints...")
-    checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)
+    #import re
+    print(" [*] Reading checkpoints...", checkpoint_dir)
+    # checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)
+    # print("     ->", checkpoint_dir)
 
     ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
     if ckpt and ckpt.model_checkpoint_path:
       ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
       self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
-      counter = int(next(re.finditer("(\d+)(?!.*\d)",ckpt_name)).group(0))
+      #counter = int(next(re.finditer("(\d+)(?!.*\d)",ckpt_name)).group(0))
+      counter = int(ckpt_name.split('-')[-1])
       print(" [*] Success to read {}".format(ckpt_name))
       return True, counter
     else:

utils.py

@@ -9,6 +9,9 @@
 import scipy.misc
 import cv2
 import numpy as np
+import os
+import time
+import datetime
 from time import gmtime, strftime
 from six.moves import xrange
 
@@ -19,6 +22,15 @@
 
 get_stddev = lambda x, k_h, k_w: 1/math.sqrt(k_w*k_h*x.get_shape()[-1])
 
+
+def expand_path(path):
+  return os.path.expanduser(os.path.expandvars(path))
+
+def timestamp(s='%Y%m%d.%H%M%S', ts=None):
+  if not ts: ts = time.time()
+  st = datetime.datetime.fromtimestamp(ts).strftime(s)
+  return st
+  
 def show_all_variables():
   model_vars = tf.trainable_variables()
   slim.model_analyzer.analyze_vars(model_vars, print_info=True)
@@ -174,12 +186,12 @@ def make_frame(t):
   clip = mpy.VideoClip(make_frame, duration=duration)
   clip.write_gif(fname, fps = len(images) / duration)
 
-def visualize(sess, dcgan, config, option):
+def visualize(sess, dcgan, config, option, sample_dir='samples'):
   image_frame_dim = int(math.ceil(config.batch_size**.5))
   if option == 0:
     z_sample = np.random.uniform(-0.5, 0.5, size=(config.batch_size, dcgan.z_dim))
     samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample})
-    save_images(samples, [image_frame_dim, image_frame_dim], './samples/test_%s.png' % strftime("%Y-%m-%d-%H-%M-%S", gmtime()))
+    save_images(samples, [image_frame_dim, image_frame_dim], os.path.join(sample_dir, 'test_%s.png' % strftime("%Y%m%d%H%M%S", gmtime() )))
   elif option == 1:
     values = np.arange(0, 1, 1./config.batch_size)
     for idx in xrange(dcgan.z_dim):
@@ -197,7 +209,7 @@ def visualize(sess, dcgan, config, option):
       else:
         samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample})
 
-      save_images(samples, [image_frame_dim, image_frame_dim], './samples/test_arange_%s.png' % (idx))
+      save_images(samples, [image_frame_dim, image_frame_dim], os.path.join(sample_dir, 'test_arange_%s.png' % (idx)))
   elif option == 2:
     values = np.arange(0, 1, 1./config.batch_size)
     for idx in [random.randint(0, dcgan.z_dim - 1) for _ in xrange(dcgan.z_dim)]:
@@ -220,7 +232,7 @@ def visualize(sess, dcgan, config, option):
       try:
         make_gif(samples, './samples/test_gif_%s.gif' % (idx))
       except:
-        save_images(samples, [image_frame_dim, image_frame_dim], './samples/test_%s.png' % strftime("%Y-%m-%d-%H-%M-%S", gmtime()))
+        save_images(samples, [image_frame_dim, image_frame_dim], os.path.join(sample_dir, 'test_%s.png' % strftime("%Y%m%d%H%M%S", gmtime() )))
   elif option == 3:
     values = np.arange(0, 1, 1./config.batch_size)
     for idx in xrange(dcgan.z_dim):
@@ -230,7 +242,7 @@ def visualize(sess, dcgan, config, option):
         z[idx] = values[kdx]
 
       samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample})
-      make_gif(samples, './samples/test_gif_%s.gif' % (idx))
+      make_gif(samples, os.path.join(sample_dir, 'test_gif_%s.gif' % (idx)))
   elif option == 4:
     image_set = []
     values = np.arange(0, 1, 1./config.batch_size)
@@ -241,7 +253,7 @@ def visualize(sess, dcgan, config, option):
       for kdx, z in enumerate(z_sample): z[idx] = values[kdx]
 
       image_set.append(sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample}))
-      make_gif(image_set[-1], './samples/test_gif_%s.gif' % (idx))
+      make_gif(image_set[-1], os.path.join(sample_dir, 'test_gif_%s.gif' % (idx)))
 
     new_image_set = [merge(np.array([images[idx] for images in image_set]), [10, 10]) \
         for idx in range(64) + range(63, -1, -1)]