DCGAN

README.md

@@ -21,6 +21,7 @@ Tensorflow implementation of [Deep Convolutional Generative Adversarial Networks
 - [Tensorflow 0.12.1](https://github.com/tensorflow/tensorflow/tree/r0.12)
 - [SciPy](http://www.scipy.org/install.html)
 - [pillow](https://github.com/python-pillow/Pillow)
+- [tqdm](https://pypi.org/project/tqdm/)
 - (Optional) [moviepy](https://github.com/Zulko/moviepy) (for visualization)
 - (Optional) [Align&Cropped Images.zip](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) : Large-scale CelebFaces Dataset
 
@@ -29,26 +30,34 @@ Tensorflow implementation of [Deep Convolutional Generative Adversarial Networks
 
 First, download dataset with:
 
-    $ python download.py --datasets mnist celebA
+    $ python download.py mnist celebA
 
 To train a model with downloaded dataset:
 
-    $ python main.py --dataset mnist --input_height=28 --output_height=28 --c_dim=1 --is_train
-    $ python main.py --dataset celebA --input_height=108 --is_train --is_crop True
+    $ python main.py --dataset mnist --input_height=28 --output_height=28 --train
+    $ python main.py --dataset celebA --input_height=108 --train --crop
 
 To test with an existing model:
 
-    $ python main.py --dataset mnist --input_height=28 --output_height=28 --c_dim=1
-    $ python main.py --dataset celebA --input_height=108 --is_crop True
+    $ python main.py --dataset mnist --input_height=28 --output_height=28
+    $ python main.py --dataset celebA --input_height=108 --crop
 
 Or, you can use your own dataset (without central crop) by:
 
     $ mkdir data/DATASET_NAME
     ... add images to data/DATASET_NAME ...
-    $ python main.py --dataset DATASET_NAME --is_train
+    $ python main.py --dataset DATASET_NAME --train
     $ python main.py --dataset DATASET_NAME
     $ # example
-    $ python main.py --dataset=eyes --input_fname_pattern="*_cropped.png" --c_dim=1 --is_train
+    $ python main.py --dataset=eyes --input_fname_pattern="*_cropped.png" --train
+
+If your dataset is located in a different root directory:
+
+    $ python main.py --dataset DATASET_NAME --data_dir DATASET_ROOT_DIR --train
+    $ python main.py --dataset DATASET_NAME --data_dir DATASET_ROOT_DIR
+    $ # example
+    $ python main.py --dataset=eyes --data_dir ../datasets/ --input_fname_pattern="*_cropped.png" --train
+
 
 ## Results
 
@@ -100,6 +109,13 @@ Details of the histogram of true and fake result of discriminator (with custom d
 ![d__hist](assets/d__hist.png)
 
 
+## Related works
+
+- [BEGAN-tensorflow](https://github.com/carpedm20/BEGAN-tensorflow)
+- [DiscoGAN-pytorch](https://github.com/carpedm20/DiscoGAN-pytorch)
+- [simulated-unsupervised-tensorflow](https://github.com/carpedm20/simulated-unsupervised-tensorflow)
+
+
 ## Author
 
 Taehoon Kim / [@carpedm20](http://carpedm20.github.io/)

download.py

@@ -172,7 +172,7 @@ def prepare_data_dir(path = './data'):
   args = parser.parse_args()
   prepare_data_dir()
 
-  if 'celebA' in args.datasets:
+  if any(name in args.datasets for name in ['CelebA', 'celebA', 'celebA']):
     download_celeb_a('./data')
   if 'lsun' in args.datasets:
     download_lsun('./data')

main.py

@@ -1,44 +1,75 @@
 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
 
 flags = tf.app.flags
 flags.DEFINE_integer("epoch", 25, "Epoch to train [25]")
 flags.DEFINE_float("learning_rate", 0.0002, "Learning rate of for adam [0.0002]")
 flags.DEFINE_float("beta1", 0.5, "Momentum term of adam [0.5]")
-flags.DEFINE_integer("train_size", np.inf, "The size of train images [np.inf]")
+flags.DEFINE_float("train_size", np.inf, "The size of train images [np.inf]")
 flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
 flags.DEFINE_integer("input_height", 108, "The size of image to use (will be center cropped). [108]")
 flags.DEFINE_integer("input_width", None, "The size of image to use (will be center cropped). If None, same value as input_height [None]")
 flags.DEFINE_integer("output_height", 64, "The size of the output images to produce [64]")
 flags.DEFINE_integer("output_width", None, "The size of the output images to produce. If None, same value as output_height [None]")
-flags.DEFINE_integer("c_dim", 3, "Dimension of image color. [3]")
 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("sample_dir", "samples", "Directory name to save the image samples [samples]")
-flags.DEFINE_boolean("is_train", False, "True for training, False for testing [False]")
-flags.DEFINE_boolean("is_crop", False, "True for training, False for testing [False]")
+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_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()
@@ -55,12 +86,15 @@ def main(_):
           batch_size=FLAGS.batch_size,
           sample_num=FLAGS.batch_size,
           y_dim=10,
-          c_dim=1,
+          z_dim=FLAGS.z_dim,
           dataset_name=FLAGS.dataset,
           input_fname_pattern=FLAGS.input_fname_pattern,
-          is_crop=FLAGS.is_crop,
+          crop=FLAGS.crop,
           checkpoint_dir=FLAGS.checkpoint_dir,
-          sample_dir=FLAGS.sample_dir)
+          sample_dir=FLAGS.sample_dir,
+          data_dir=FLAGS.data_dir,
+          out_dir=FLAGS.out_dir,
+          max_to_keep=FLAGS.max_to_keep)
     else:
       dcgan = DCGAN(
           sess,
@@ -70,20 +104,25 @@ def main(_):
           output_height=FLAGS.output_height,
           batch_size=FLAGS.batch_size,
           sample_num=FLAGS.batch_size,
-          c_dim=FLAGS.c_dim,
+          z_dim=FLAGS.z_dim,
           dataset_name=FLAGS.dataset,
           input_fname_pattern=FLAGS.input_fname_pattern,
-          is_crop=FLAGS.is_crop,
+          crop=FLAGS.crop,
           checkpoint_dir=FLAGS.checkpoint_dir,
-          sample_dir=FLAGS.sample_dir)
+          sample_dir=FLAGS.sample_dir,
+          data_dir=FLAGS.data_dir,
+          out_dir=FLAGS.out_dir,
+          max_to_keep=FLAGS.max_to_keep)
 
     show_all_variables()
-    if FLAGS.is_train:
+
+    if FLAGS.train:
       dcgan.train(FLAGS)
     else:
-      if not dcgan.load(FLAGS.checkpoint_dir):
-        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],
@@ -92,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()