From 2463113f77fbf382cb8b18835ebc383f4bc16f5d Mon Sep 17 00:00:00 2001
From: Wesley Emeneker <1059409812@mil>
Date: Wed, 15 Aug 2018 16:45:53 -1000
Subject: [PATCH 1/7] Fixing some basic PEP8 stuff. Getting familiar with code.

End goal is multiGPU support via batch-splitting for data parallelization.
---
 networks/models/dcgan.py | 50 ++++++++++++++++++++--------------------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/networks/models/dcgan.py b/networks/models/dcgan.py
index b90d1a2..d56c567 100644
--- a/networks/models/dcgan.py
+++ b/networks/models/dcgan.py
@@ -1,9 +1,10 @@
 import tensorflow as tf
 from .ops import linear, conv2d, conv2d_transpose, lrelu
 
-class dcgan(object):
-    def __init__(self, output_size=64, batch_size=64, 
-                 nd_layers=4, ng_layers=4, df_dim=128, gf_dim=128, 
+
+class DCGAN(object):
+    def __init__(self, output_size=64, batch_size=64,
+                 nd_layers=4, ng_layers=4, df_dim=128, gf_dim=128,
                  c_dim=1, z_dim=100, flip_labels=0.01, data_format="NHWC",
                  gen_prior=tf.random_normal, transpose_b=False):
 
@@ -18,13 +19,12 @@ def __init__(self, output_size=64, batch_size=64,
         self.flip_labels = flip_labels
         self.data_format = data_format
         self.gen_prior = gen_prior
-        self.transpose_b = transpose_b # transpose weight matrix in linear layers for (possible) better performance when running on HSW/KNL
-        self.stride = 2 # this is fixed for this architecture
+        self.transpose_b = transpose_b  # transpose weight matrix in linear layers for (possible) better performance when running on HSW/KNL
+        self.stride = 2  # this is fixed for this architecture
 
         self._check_architecture_consistency()
 
-
-        self.batchnorm_kwargs = {'epsilon' : 1e-5, 'decay': 0.9, 
+        self.batchnorm_kwargs = {'epsilon': 1e-5, 'decay': 0.9,
                                  'updates_collections': None, 'scale': True,
                                  'fused': True, 'data_format': self.data_format}
 
@@ -37,10 +37,10 @@ def training_graph(self):
 
         self.z = self.gen_prior(shape=[self.batch_size, self.z_dim])
 
-        with tf.variable_scope("discriminator") as d_scope:
+        with tf.variable_scope("discriminator"):
             d_prob_real, d_logits_real = self.discriminator(self.images, is_training=True)
 
-        with tf.variable_scope("generator") as g_scope:
+        with tf.variable_scope("generator"):
             g_images = self.generator(self.z, is_training=True)
 
         with tf.variable_scope("discriminator") as d_scope:
@@ -63,7 +63,7 @@ def training_graph(self):
                                            tf.summary.scalar("loss/d", self.d_loss)])
 
         g_sum = [tf.summary.scalar("loss/g", self.g_loss)]
-        if self.data_format == "NHWC": # tf.summary.image is not implemented for NCHW
+        if self.data_format == "NHWC":  # tf.summary.image is not implemented for NCHW
             g_sum.append(tf.summary.image("G", g_images, max_outputs=4))
         self.g_summary = tf.summary.merge(g_sum)
 
@@ -71,7 +71,7 @@ def training_graph(self):
         self.d_vars = [var for var in t_vars if 'discriminator/' in var.name]
         self.g_vars = [var for var in t_vars if 'generator/' in var.name]
 
-        with tf.variable_scope("counters") as counters_scope:
+        with tf.variable_scope("counters"):
             self.epoch = tf.Variable(-1, name='epoch', trainable=False)
             self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
             self.global_step = tf.train.get_or_create_global_step()
@@ -87,13 +87,13 @@ def inference_graph(self):
 
         self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
 
-        with tf.variable_scope("discriminator") as d_scope:
-            self.D,_ = self.discriminator(self.images, is_training=False)
+        with tf.variable_scope("discriminator"):
+            self.D, _ = self.discriminator(self.images, is_training=False)
 
-        with tf.variable_scope("generator") as g_scope:
+        with tf.variable_scope("generator"):
             self.G = self.generator(self.z, is_training=False)
 
-        with tf.variable_scope("counters") as counters_scope:
+        with tf.variable_scope("counters"):
             self.epoch = tf.Variable(-1, name='epoch', trainable=False)
             self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
             self.global_step = tf.train.get_or_create_global_step()
@@ -110,11 +110,9 @@ def optimizer(self, learning_rate, beta1):
 
         return tf.group(d_optim, g_optim, name="all_optims")
 
-                                                   
     def generator(self, z, is_training):
-
         map_size = self.output_size/int(2**self.ng_layers)
-        num_channels = self.gf_dim * int(2**(self.ng_layers -1))
+        num_channels = self.gf_dim * int(2**(self.ng_layers - 1))
 
         # h0 = relu(BN(reshape(FC(z))))
         z_ = linear(z, num_channels*map_size*map_size, 'h0_lin', transpose_b=self.transpose_b)
@@ -129,19 +127,18 @@ def generator(self, z, is_training):
             num_channels /= 2
             chain = conv2d_transpose(chain,
                                      self._tensor_data_format(self.batch_size, map_size, map_size, num_channels),
-                                     stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T'%h)
-            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i'%h, **self.batchnorm_kwargs)
+                                     stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % h)
+            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
             chain = tf.nn.relu(chain)
 
         # h1 = conv2d_transpose(h0)
         map_size *= self.stride
         hn = conv2d_transpose(chain,
                               self._tensor_data_format(self.batch_size, map_size, map_size, self.c_dim),
-                              stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T'%(self.ng_layers))
+                              stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % self.ng_layers)
 
         return tf.nn.tanh(hn)
 
-
     def discriminator(self, image, is_training):
 
         # h0 = lrelu(conv2d(image))
@@ -150,12 +147,15 @@ def discriminator(self, image, is_training):
         chain = h0
         for h in range(1, self.nd_layers):
             # h1 = lrelu(BN(conv2d(h0)))
-            chain = conv2d(chain, self.df_dim*(2**h), self.data_format, name='h%i_conv'%h)
-            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i'%h, **self.batchnorm_kwargs)
+            chain = conv2d(chain, self.df_dim*(2**h), self.data_format, name='h%i_conv' % h)
+            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
             chain = lrelu(chain)
 
         # h1 = linear(reshape(h0))
-        hn = linear(tf.reshape(chain, [self.batch_size, -1]), 1, 'h%i_lin'%self.nd_layers, transpose_b=self.transpose_b)
+        hn = linear(tf.reshape(chain, [self.batch_size, -1]),
+                    1,
+                    'h%i_lin' % self.nd_layers,
+                    transpose_b=self.transpose_b)
 
         return tf.nn.sigmoid(hn), hn
 

From db7e1b53d10654e577c7d6df77f805684dbb3b1b Mon Sep 17 00:00:00 2001
From: Wesley Emeneker <wesley.d.emeneker.ctr@mhpcc.hpc.mil>
Date: Mon, 20 Aug 2018 15:26:38 -1000
Subject: [PATCH 2/7] Still working on multigpu split. Not working yet.

In the midst of understanding the code and averaging the split gradients
  so that the generator and discriminator learn the same things.
On the first round of training, some of the generator variables
 (h0_lin, biases, etc.) have None for gradients.
I'm not sure why, or if that will change.
---
 networks/models/dcgan.py | 436 ++++++++++++++++++++++-----------------
 networks/models/main.py  |  93 +++++----
 networks/models/ops.py   | 196 ++++++++++++------
 networks/models/train.py | 153 +++++++-------
 networks/models/utils.py |  82 ++++----
 networks/run_dcgan.py    |  74 +++----
 6 files changed, 580 insertions(+), 454 deletions(-)

diff --git a/networks/models/dcgan.py b/networks/models/dcgan.py
index d56c567..302b612 100644
--- a/networks/models/dcgan.py
+++ b/networks/models/dcgan.py
@@ -1,193 +1,243 @@
-import tensorflow as tf
-from .ops import linear, conv2d, conv2d_transpose, lrelu
-
-
-class DCGAN(object):
-    def __init__(self, output_size=64, batch_size=64,
-                 nd_layers=4, ng_layers=4, df_dim=128, gf_dim=128,
-                 c_dim=1, z_dim=100, flip_labels=0.01, data_format="NHWC",
-                 gen_prior=tf.random_normal, transpose_b=False):
-
-        self.output_size = output_size
-        self.batch_size = batch_size
-        self.nd_layers = nd_layers
-        self.ng_layers = ng_layers
-        self.df_dim = df_dim
-        self.gf_dim = gf_dim
-        self.c_dim = c_dim
-        self.z_dim = z_dim
-        self.flip_labels = flip_labels
-        self.data_format = data_format
-        self.gen_prior = gen_prior
-        self.transpose_b = transpose_b  # transpose weight matrix in linear layers for (possible) better performance when running on HSW/KNL
-        self.stride = 2  # this is fixed for this architecture
-
-        self._check_architecture_consistency()
-
-        self.batchnorm_kwargs = {'epsilon': 1e-5, 'decay': 0.9,
-                                 'updates_collections': None, 'scale': True,
-                                 'fused': True, 'data_format': self.data_format}
-
-    def training_graph(self):
-
-        if self.data_format == "NHWC":
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
-        else:
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
-
-        self.z = self.gen_prior(shape=[self.batch_size, self.z_dim])
-
-        with tf.variable_scope("discriminator"):
-            d_prob_real, d_logits_real = self.discriminator(self.images, is_training=True)
-
-        with tf.variable_scope("generator"):
-            g_images = self.generator(self.z, is_training=True)
-
-        with tf.variable_scope("discriminator") as d_scope:
-            d_scope.reuse_variables()
-            d_prob_fake, d_logits_fake = self.discriminator(g_images, is_training=True)
-
-        with tf.name_scope("losses"):
-            with tf.name_scope("d"):
-                d_label_real, d_label_fake = self._labels()
-                self.d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real, labels=d_label_real, name="real"))
-                self.d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=d_label_fake, name="fake"))
-                self.d_loss = self.d_loss_real + self.d_loss_fake
-            with tf.name_scope("g"):
-                self.g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_logits_fake)))
-
-        self.d_summary = tf.summary.merge([tf.summary.histogram("prob/real", d_prob_real),
-                                           tf.summary.histogram("prob/fake", d_prob_fake),
-                                           tf.summary.scalar("loss/real", self.d_loss_real),
-                                           tf.summary.scalar("loss/fake", self.d_loss_fake),
-                                           tf.summary.scalar("loss/d", self.d_loss)])
-
-        g_sum = [tf.summary.scalar("loss/g", self.g_loss)]
-        if self.data_format == "NHWC":  # tf.summary.image is not implemented for NCHW
-            g_sum.append(tf.summary.image("G", g_images, max_outputs=4))
-        self.g_summary = tf.summary.merge(g_sum)
-
-        t_vars = tf.trainable_variables()
-        self.d_vars = [var for var in t_vars if 'discriminator/' in var.name]
-        self.g_vars = [var for var in t_vars if 'generator/' in var.name]
-
-        with tf.variable_scope("counters"):
-            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
-            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
-            self.global_step = tf.train.get_or_create_global_step()
-
-        self.saver = tf.train.Saver(max_to_keep=8000)
-
-    def inference_graph(self):
-
-        if self.data_format == "NHWC":
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
-        else:
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
-
-        self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
-
-        with tf.variable_scope("discriminator"):
-            self.D, _ = self.discriminator(self.images, is_training=False)
-
-        with tf.variable_scope("generator"):
-            self.G = self.generator(self.z, is_training=False)
-
-        with tf.variable_scope("counters"):
-            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
-            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
-            self.global_step = tf.train.get_or_create_global_step()
-
-        self.saver = tf.train.Saver(max_to_keep=8000)
-
-    def optimizer(self, learning_rate, beta1):
-
-        d_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1) \
-                                         .minimize(self.d_loss, var_list=self.d_vars, global_step=self.global_step)
-
-        g_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1) \
-                                         .minimize(self.g_loss, var_list=self.g_vars)
-
-        return tf.group(d_optim, g_optim, name="all_optims")
-
-    def generator(self, z, is_training):
-        map_size = self.output_size/int(2**self.ng_layers)
-        num_channels = self.gf_dim * int(2**(self.ng_layers - 1))
-
-        # h0 = relu(BN(reshape(FC(z))))
-        z_ = linear(z, num_channels*map_size*map_size, 'h0_lin', transpose_b=self.transpose_b)
-        h0 = tf.reshape(z_, self._tensor_data_format(-1, map_size, map_size, num_channels))
-        bn0 = tf.contrib.layers.batch_norm(h0, is_training=is_training, scope='bn0', **self.batchnorm_kwargs)
-        h0 = tf.nn.relu(bn0)
-
-        chain = h0
-        for h in range(1, self.ng_layers):
-            # h1 = relu(BN(conv2d_transpose(h0)))
-            map_size *= self.stride
-            num_channels /= 2
-            chain = conv2d_transpose(chain,
-                                     self._tensor_data_format(self.batch_size, map_size, map_size, num_channels),
-                                     stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % h)
-            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
-            chain = tf.nn.relu(chain)
-
-        # h1 = conv2d_transpose(h0)
-        map_size *= self.stride
-        hn = conv2d_transpose(chain,
-                              self._tensor_data_format(self.batch_size, map_size, map_size, self.c_dim),
-                              stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % self.ng_layers)
-
-        return tf.nn.tanh(hn)
-
-    def discriminator(self, image, is_training):
-
-        # h0 = lrelu(conv2d(image))
-        h0 = lrelu(conv2d(image, self.df_dim, self.data_format, name='h0_conv'))
-
-        chain = h0
-        for h in range(1, self.nd_layers):
-            # h1 = lrelu(BN(conv2d(h0)))
-            chain = conv2d(chain, self.df_dim*(2**h), self.data_format, name='h%i_conv' % h)
-            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
-            chain = lrelu(chain)
-
-        # h1 = linear(reshape(h0))
-        hn = linear(tf.reshape(chain, [self.batch_size, -1]),
-                    1,
-                    'h%i_lin' % self.nd_layers,
-                    transpose_b=self.transpose_b)
-
-        return tf.nn.sigmoid(hn), hn
-
-    def _labels(self):
-        with tf.name_scope("labels"):
-            ones = tf.ones([self.batch_size, 1])
-            zeros = tf.zeros([self.batch_size, 1])
-            flip_labels = tf.constant(self.flip_labels)
-
-            if self.flip_labels > 0:
-                prob = tf.random_uniform([], 0, 1)
-
-                d_label_real = tf.cond(tf.less(prob, flip_labels), lambda: zeros, lambda: ones)
-                d_label_fake = tf.cond(tf.less(prob, flip_labels), lambda: ones, lambda: zeros)
-            else:
-                d_label_real = ones
-                d_label_fake = zeros
-
-        return d_label_real, d_label_fake
-
-    def _tensor_data_format(self, N, H, W, C):
-        if self.data_format == "NHWC":
-            return [int(N), int(H), int(W), int(C)]
-        else:
-            return [int(N), int(C), int(H), int(W)]
-
-    def _check_architecture_consistency(self):
-
-        if self.output_size/2**self.nd_layers < 1:
-            print("Error: Number of discriminator conv. layers are larger than the output_size for this architecture")
-            exit(0)
-
-        if self.output_size/2**self.ng_layers < 1:
-            print("Error: Number of generator conv_transpose layers are larger than the output_size for this architecture")
-            exit(0)
+import sys
+
+import tensorflow as tf
+from .ops import linear, conv2d, conv2d_transpose, lrelu, average_gradients, get_available_gpus
+
+
+class DCGAN(object):
+    def __init__(self, output_size=64, batch_size=64,
+                 nd_layers=4, ng_layers=4, df_dim=128, gf_dim=128,
+                 c_dim=1, z_dim=100, flip_labels=0.01, data_format="NHWC",
+                 gen_prior=tf.random_normal, transpose_b=False,
+                 num_gpus=-1, learning_rate=.0002, beta1=0.5):
+
+        self.output_size = output_size
+        self.batch_size = batch_size
+        self.nd_layers = nd_layers # Number of hidden layers in the discriminator?
+        self.ng_layers = ng_layers # Number of hidden layers in the generator?
+        self.df_dim = df_dim # discriminator f dimensions?
+        self.gf_dim = gf_dim # generator f dimensions
+        self.c_dim = c_dim # number of channels
+        self.z_dim = z_dim # 
+        self.flip_labels = flip_labels
+        self.data_format = data_format
+        self.gen_prior = gen_prior
+        self.transpose_b = transpose_b  # transpose weight matrix in linear layers for (possible) better performance when running on HSW/KNL
+        self.stride = 2  # this is fixed for this architecture
+        
+        self.compute_devices = get_available_gpus()
+        # Figure out if we are using gpus, how many, which ones
+        # If num_gpus < 0, use all GPUs we were given
+        if num_gpus > 0:
+            # If the user doesn't want to use all the gpus, take the first N
+            self.compute_devices = self.compute_devices[:num_gpus]
+        elif num_gpus == 0:
+            self.compute_devices = ['/cpu:0']
+
+        self._check_architecture_consistency()
+
+        self.batchnorm_kwargs = {'epsilon': 1e-5, 'decay': 0.9,
+                                 'updates_collections': None, 'scale': True,
+                                 'fused': True, 'data_format': self.data_format}
+        self.make_optimizers(learning_rate, beta1)
+
+    @property
+    def compute_batch_size(self):
+        """Return the batch size that each GPU (or other compute device) should have"""
+        return self.batch_size / len(self.compute_devices)
+        
+    def training_graph(self):
+        """Make the training graph such that it divides the work among 1 or more GPUS"""
+
+        if self.data_format == "NHWC":
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
+        else:
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
+            
+        self.z = self.gen_prior(shape=[self.batch_size, self.z_dim])
+
+        # Take the batch of images and split the generator and discriminator work among
+        #   all the computing devices given.
+        # My hope is that by parallelizing this inner graph, we'll get some
+        #   nice speedup without having to refactor a lot.
+        real_image_splits = tf.split(self.images, num_or_size_splits=len(self.compute_devices), axis=0)
+        # Split the random-noise images into chunks, one for each compute device
+        gen_image_splits = tf.split(self.z, num_or_size_splits=len(self.compute_devices), axis=0)
+
+        tower_d_gradients = []
+        tower_g_gradients = []
+        with tf.variable_scope(tf.get_variable_scope()):
+            for idx, (images, z_images) in enumerate(zip(real_image_splits, gen_image_splits)):
+                compute_device = self.compute_devices[idx]
+                with tf.device(compute_device):
+                    with tf.variable_scope("tower_%d" % idx) as scope:
+                        with tf.variable_scope("discriminator"):
+                            d_prob_real, d_logits_real = self.discriminator(images, is_training=True)
+
+                        with tf.variable_scope("generator"):
+                            g_images = self.generator(z_images, is_training=True)
+
+                        print >>sys.stderr, "SHAPES:", g_images.shape, images.shape, z_images.shape, self.batch_size, self.compute_devices
+                        with tf.variable_scope("discriminator") as d_scope:
+                            d_scope.reuse_variables()
+                            d_prob_fake, d_logits_fake = self.discriminator(g_images, is_training=True)
+
+                        with tf.name_scope("losses"):
+                            with tf.name_scope("d"):
+                                d_label_real, d_label_fake = self._labels()
+                                d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real, labels=d_label_real, name="real"))
+                                d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=d_label_fake, name="fake"))
+                                d_loss = d_loss_real + d_loss_fake
+
+                            with tf.name_scope("g"):
+                                g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_logits_fake)))
+
+                        t_vars = tf.trainable_variables()
+                        d_vars = [var for var in t_vars if 'discriminator/' in var.name and "gradients" in var.name]
+                        g_vars = [var for var in t_vars if 'generator/' in var.name and "gradients" in var.names]
+
+                        d_gradients, g_gradients = self.compute_gradients(d_loss, g_loss, d_vars, g_vars)
+                        tower_d_gradients.append(d_gradients)
+                        tower_g_gradients.append(g_gradients)
+
+        # Do these variables need to be separated by variable name so that
+        #   the average_gradients function will correctly aggregate them?
+        # I think so.
+        g_grads_vars = average_gradients(tower_g_gradients)
+        d_grads_vars = average_gradients(tower_d_gradients)
+                        
+        self.d_summary = tf.summary.merge([tf.summary.histogram("prob/real", d_prob_real),
+                                           tf.summary.histogram("prob/fake", d_prob_fake),
+                                           tf.summary.scalar("loss/real", d_loss_real),
+                                           tf.summary.scalar("loss/fake", d_loss_fake),
+                                           tf.summary.scalar("loss/d", d_loss)])
+
+        g_sum = [tf.summary.scalar("loss/g", g_loss)]
+        if self.data_format == "NHWC":  # tf.summary.image is not implemented for NCHW
+            g_sum.append(tf.summary.image("G", g_images, max_outputs=4))
+        self.g_summary = tf.summary.merge(g_sum)
+
+        with tf.variable_scope("counters"):
+            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
+            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
+            self.global_step = tf.train.get_or_create_global_step()
+
+        self.saver = tf.train.Saver(max_to_keep=8000)
+
+        return g_grads_vars, d_grads_vars
+        
+    def inference_graph(self):
+
+        if self.data_format == "NHWC":
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
+        else:
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
+
+        self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
+
+        with tf.variable_scope("discriminator"):
+            self.D, _ = self.discriminator(self.images, is_training=False)
+
+        with tf.variable_scope("generator"):
+            self.G = self.generator(self.z, is_training=False)
+
+        with tf.variable_scope("counters"):
+            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
+            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
+            self.global_step = tf.train.get_or_create_global_step()
+
+        self.saver = tf.train.Saver(max_to_keep=8000)
+
+    def make_optimizers(self, learning_rate, beta1):
+        self.d_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1)
+        self.g_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1)
+        
+    def compute_gradients(self, d_loss, g_loss, d_vars, g_vars):
+        d_optim = self.d_optim.compute_gradients(d_loss, var_list=d_vars)
+        g_optim = self.g_optim.compute_gradients(g_loss, var_list=g_vars)
+
+        return d_optim, g_optim
+
+    def apply_gradients(g_grads_vars, d_grads_vars):
+        return tf.group(self.g_optim.apply_gradients(g_grads_vars),
+                        self.d_optim.apply_gradients(d_grads_vars, global_step=self.global_step))
+
+    def generator(self, z, is_training):
+        map_size = self.output_size / int(2**self.ng_layers)
+        num_channels = self.gf_dim * int(2**(self.ng_layers - 1))
+
+        z_ = linear(z, num_channels*map_size*map_size, 'h0_lin', transpose_b=self.transpose_b)
+        h0 = tf.reshape(z_, self._tensor_data_format(-1, map_size, map_size, num_channels))
+        bn0 = tf.contrib.layers.batch_norm(h0, is_training=is_training, scope='bn0', **self.batchnorm_kwargs)
+        h0 = tf.nn.relu(bn0)
+
+        chain = h0
+
+        print >>sys.stderr, "chain shape:", chain.shape
+        
+        for h in range(1, self.ng_layers):
+            map_size *= self.stride
+            num_channels /= 2
+            chain = conv2d_transpose(chain,
+                                     self._tensor_data_format(self.compute_batch_size, map_size, map_size, num_channels),
+                                     stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % h)
+            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
+            chain = tf.nn.relu(chain)
+
+        map_size *= self.stride
+        hn = conv2d_transpose(chain,
+                              self._tensor_data_format(self.compute_batch_size, map_size, map_size, self.c_dim),
+                              stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % self.ng_layers)
+
+        return tf.nn.tanh(hn)
+
+    def discriminator(self, image, is_training):
+        """This discriminator works on a single image at a time"""
+
+        chain = lrelu(conv2d(image, self.df_dim, self.data_format, name='h0_conv'))
+
+        for h in range(1, self.nd_layers):
+            chain = conv2d(chain, self.df_dim*(2**h), self.data_format, name='h%i_conv' % h)
+            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
+            chain = lrelu(chain)
+
+        hn = linear(tf.reshape(chain, [self.batch_size, -1]),
+                    1,
+                    'h%i_lin' % self.nd_layers,
+                    transpose_b=self.transpose_b)
+
+        return tf.nn.sigmoid(hn), hn
+
+    def _labels(self):
+        with tf.name_scope("labels"):
+            ones = tf.ones([self.batch_size, 1])
+            zeros = tf.zeros([self.batch_size, 1])
+            flip_labels = tf.constant(self.flip_labels)
+
+            if self.flip_labels > 0:
+                prob = tf.random_uniform([], 0, 1)
+
+                d_label_real = tf.cond(tf.less(prob, flip_labels), lambda: zeros, lambda: ones)
+                d_label_fake = tf.cond(tf.less(prob, flip_labels), lambda: ones, lambda: zeros)
+            else:
+                d_label_real = ones
+                d_label_fake = zeros
+
+        return d_label_real, d_label_fake
+
+    def _tensor_data_format(self, N, H, W, C):
+        if self.data_format == "NHWC":
+            return [int(N), int(H), int(W), int(C)]
+        else:
+            return [int(N), int(C), int(H), int(W)]
+
+    def _check_architecture_consistency(self):
+
+        if self.output_size/2**self.nd_layers < 1:
+            print("Error: Number of discriminator conv. layers are larger than the output_size for this architecture")
+            exit(0)
+
+        if self.output_size/2**self.ng_layers < 1:
+            print("Error: Number of generator conv_transpose layers are larger than the output_size for this architecture")
+            exit(0)
+
diff --git a/networks/models/main.py b/networks/models/main.py
index 30b9bed..af9c676 100644
--- a/networks/models/main.py
+++ b/networks/models/main.py
@@ -1,45 +1,48 @@
-import tensorflow as tf
-import train
-import numpy as np
-import pprint
-
-flags = tf.app.flags
-flags.DEFINE_string("dataset", "cosmo", "The name of dataset [cosmo]")
-flags.DEFINE_string("datafile", "data/cosmo_primary_64_1k_train.npy", "Input data file for cosmo")
-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_float("flip_labels", 0, "Probability of flipping labels [0]")
-flags.DEFINE_integer("z_dim", 100, "Dimension of noise vector z [100]")
-flags.DEFINE_integer("nd_layers", 4, "Number of discriminator convolutional layers. [4]")
-flags.DEFINE_integer("ng_layers", 4, "Number of generator conv_T layers. [4]")
-flags.DEFINE_integer("gf_dim", 64, "Dimension of gen filters in last conv layer. [64]")
-flags.DEFINE_integer("df_dim", 64, "Dimension of discrim filters in first conv layer. [64]")
-flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
-flags.DEFINE_integer("output_size", 64, "The size of the output images to produce [64]")
-flags.DEFINE_integer("c_dim", 1, "Dimension of image color. [1]")
-flags.DEFINE_string("data_format", "NHWC", "data format [NHWC]")
-flags.DEFINE_boolean("transpose_matmul_b", False, "Transpose matmul B matrix for performance [False]")
-flags.DEFINE_string("checkpoint_dir", "checkpoints", "Directory name to save the checkpoints [checkpoint]")
-flags.DEFINE_string("experiment", "run_0", "Tensorboard run directory name [run_0]")
-flags.DEFINE_boolean("save_every_step", False, "Save a checkpoint after every step [False]")
-flags.DEFINE_boolean("verbose", True, "print loss on every step [False]")
-config = flags.FLAGS
-
-def main(_):
-
-    pprint.PrettyPrinter().pprint(config.__flags)
-    train.train_dcgan(get_data(), config)
-
-def get_data():
-    data = np.load(config.datafile, mmap_mode='r')
-
-    if config.data_format == 'NHWC':
-        data = np.expand_dims(data, axis=-1)
-    else: # 'NCHW'
-        data = np.expand_dims(data, axis=1)
-
-    return data
-
-if __name__ == '__main__':
-    tf.app.run()
+import tensorflow as tf
+import train
+import numpy as np
+import pprint
+
+flags = tf.app.flags
+flags.DEFINE_string("dataset", "cosmo", "The name of dataset [cosmo]")
+flags.DEFINE_string("datafile", "data/cosmo_primary_64_1k_train.npy", "Input data file for cosmo")
+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_float("flip_labels", 0, "Probability of flipping labels [0]")
+flags.DEFINE_integer("z_dim", 100, "Dimension of noise vector z [100]")
+flags.DEFINE_integer("nd_layers", 4, "Number of discriminator convolutional layers. [4]")
+flags.DEFINE_integer("ng_layers", 4, "Number of generator conv_T layers. [4]")
+flags.DEFINE_integer("gf_dim", 64, "Dimension of gen filters in last conv layer. [64]")
+flags.DEFINE_integer("df_dim", 64, "Dimension of discrim filters in first conv layer. [64]")
+flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
+flags.DEFINE_integer("output_size", 64, "The size of the output images to produce [64]")
+flags.DEFINE_integer("c_dim", 1, "Dimension of image color. [1]")
+flags.DEFINE_string("data_format", "NHWC", "data format [NHWC]")
+flags.DEFINE_boolean("transpose_matmul_b", False, "Transpose matmul B matrix for performance [False]")
+flags.DEFINE_string("checkpoint_dir", "checkpoints", "Directory name to save the checkpoints [checkpoint]")
+flags.DEFINE_string("experiment", "run_0", "Tensorboard run directory name [run_0]")
+flags.DEFINE_boolean("save_every_step", False, "Save a checkpoint after every step [False]")
+flags.DEFINE_boolean("verbose", True, "print loss on every step [False]")
+flags.DEFINE_integer("num_gpus", -1, "Number of GPUs to use to train GAN. [-1]. -1 means use all available GPUs.")
+config = flags.FLAGS
+
+
+def main(_):
+    pprint.PrettyPrinter().pprint(config.__flags)
+    train.train_dcgan(get_data(), config)
+
+
+def get_data():
+    data = np.load(config.datafile, mmap_mode='r')
+
+    if config.data_format == 'NHWC':
+        data = np.expand_dims(data, axis=-1)
+    else:  # 'NCHW'
+        data = np.expand_dims(data, axis=1)
+
+    return data
+
+
+if __name__ == '__main__':
+    tf.app.run()
diff --git a/networks/models/ops.py b/networks/models/ops.py
index ca79d93..2d1ddc5 100644
--- a/networks/models/ops.py
+++ b/networks/models/ops.py
@@ -1,64 +1,132 @@
-import tensorflow as tf
-
-def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, transpose_b=False):
-
-    shape = input_.get_shape().as_list()
-    if not transpose_b:
-        w_shape = [shape[1], output_size]
-    else:
-        w_shape = [output_size, shape[1]]
-
-    with tf.variable_scope(scope or "linear"):
-        matrix = tf.get_variable('w', w_shape, tf.float32,
-                                 tf.random_normal_initializer(stddev=stddev))
-        bias = tf.get_variable('b', [output_size],
-            initializer=tf.constant_initializer(bias_start))
-
-        return tf.matmul(input_, matrix, transpose_b=transpose_b) + bias
-
-def conv2d(input_, out_channels, data_format, kernel=5, stride=2, stddev=0.02, name="conv2d"):
-
-    if data_format == "NHWC":
-        in_channels = input_.get_shape()[-1]
-        strides = [1, stride, stride, 1]
-    else: # NCHW
-        in_channels = input_.get_shape()[1]
-        strides = [1, 1, stride, stride]
-
-    with tf.variable_scope(name):
-        w = tf.get_variable('w', [kernel, kernel, in_channels, out_channels],
-                            initializer=tf.truncated_normal_initializer(stddev=stddev))
-        conv = tf.nn.conv2d(input_, w, strides=strides, padding='SAME', data_format=data_format)
-
-        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
-        conv = tf.reshape(tf.nn.bias_add(conv, biases, data_format=data_format), conv.get_shape())
-
-        return conv
-
-def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stddev=0.02,
-                     name="conv2d_transpose"):
-
-    if data_format == "NHWC":
-        in_channels = input_.get_shape()[-1]
-        out_channels = output_shape[-1]
-        strides = [1, stride, stride, 1]
-    else:
-        in_channels = input_.get_shape()[1]
-        out_channels = output_shape[1]
-        strides = [1, 1, stride, stride]
-
-    with tf.variable_scope(name):
-        w = tf.get_variable('w', [kernel, kernel, out_channels, in_channels],
-                            initializer=tf.random_normal_initializer(stddev=stddev))
-        
-        deconv = tf.nn.conv2d_transpose(input_, w, output_shape=output_shape, strides=strides, data_format=data_format)
-
-        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
-        deconv = tf.reshape(tf.nn.bias_add(deconv, biases, data_format=data_format), deconv.get_shape())
-
-        return deconv
-       
-
-def lrelu(x, alpha=0.2, name="lrelu"):
-    with tf.name_scope(name):
-      return tf.maximum(x, alpha*x)
+from pprint import pprint
+import sys
+
+import tensorflow as tf
+
+
+def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, transpose_b=False):
+    """Make and return a matrix multiply (input_ * w) + b(ias) operation.
+    
+    """
+
+    shape = input_.get_shape().as_list()
+    if not transpose_b:
+        w_shape = [shape[1], output_size]
+    else:
+        w_shape = [output_size, shape[1]]
+
+    with tf.variable_scope(scope or "linear"):
+        matrix = tf.get_variable('w', w_shape, tf.float32,
+                                 tf.random_normal_initializer(stddev=stddev))
+        bias = tf.get_variable('b', [output_size],
+                               initializer=tf.constant_initializer(bias_start))
+
+        return tf.matmul(input_, matrix, transpose_b=transpose_b) + bias
+
+
+def conv2d(input_, out_channels, data_format, kernel=5, stride=2, stddev=0.02, name="conv2d"):
+    """Make and return a 2D convolution + bias operation.
+    """
+    if data_format == "NHWC":
+        in_channels = input_.get_shape()[-1]
+        strides = [1, stride, stride, 1]
+    else:  # NCHW
+        # WARNING: These strides are probably broken
+        #   https://www.tensorflow.org/api_docs/python/tf/nn/conv2d:
+        #    "Must have strides[0] = strides[3] = 1"
+        in_channels = input_.get_shape()[1]
+        strides = [1, 1, stride, stride]
+
+    with tf.variable_scope(name):
+        # By default, our convolutional mask is a 5x5 filter.
+        w = tf.get_variable('w', [kernel, kernel, in_channels, out_channels],
+                            initializer=tf.truncated_normal_initializer(stddev=stddev))
+        conv = tf.nn.conv2d(input_, w, strides=strides, padding='SAME', data_format=data_format)
+
+        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
+        conv = tf.reshape(tf.nn.bias_add(conv, biases, data_format=data_format), conv.get_shape())
+
+        return conv
+
+
+def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stddev=0.02,
+                     name="conv2d_transpose"):
+
+    if data_format == "NHWC":
+        in_channels = input_.get_shape()[-1]
+        out_channels = output_shape[-1]
+        strides = [1, stride, stride, 1]
+    else:
+        # WARNING: These strides are probably broken
+        #   https://www.tensorflow.org/api_docs/python/tf/nn/conv2d:
+        #    "Must have strides[0] = strides[3] = 1"
+        in_channels = input_.get_shape()[1]
+        out_channels = output_shape[1]
+        strides = [1, 1, stride, stride]
+
+    with tf.variable_scope(name):
+        w = tf.get_variable('w', [kernel, kernel, out_channels, in_channels],
+                            initializer=tf.random_normal_initializer(stddev=stddev))
+
+        deconv = tf.nn.conv2d_transpose(input_, w, output_shape=output_shape, strides=strides, data_format=data_format)
+
+        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
+        deconv = tf.reshape(tf.nn.bias_add(deconv, biases, data_format=data_format), deconv.get_shape())
+
+        return deconv
+
+
+def lrelu(x, alpha=0.2, name="lrelu"):
+    with tf.name_scope(name):
+        return tf.maximum(x, alpha*x)
+
+
+def average_gradients(tower_grads):
+    """Taken from https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py
+
+    Calculate the average gradient for each shared variable across all towers.
+
+    Note that this function provides a synchronization point across all towers.
+    
+    Args:
+    tower_grads: List of lists of (gradient, variable) tuples. The outer list
+      is over individual gradients. The inner list is over the gradient
+      calculation for each tower.
+    Returns:
+     List of pairs of (gradient, variable) where the gradient has been averaged
+     across all towers.
+    """
+    pprint(tower_grads, sys.stderr)
+    average_grads = []
+    for grad_and_vars in zip(*tower_grads):
+        # Note that each grad_and_vars looks like the following:
+        #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
+        grads = []
+        for g, _ in grad_and_vars:
+            print >>sys.stderr, "var name:", _.name, _
+            if g is None:
+                pass
+            # Add 0 dimension to the gradients to represent the tower.
+            expanded_g = tf.expand_dims(g, 0)
+
+            # Append on a 'tower' dimension which we will average over below.
+            grads.append(expanded_g)
+
+        # Average over the 'tower' dimension.
+        grad = tf.concat(axis=0, values=grads)
+        grad = tf.reduce_mean(grad, 0)
+
+        # Keep in mind that the Variables are redundant because they are shared
+        # across towers. So .. we will just return the first tower's pointer to
+        # the Variable.
+        v = grad_and_vars[0][1]
+        grad_and_var = (grad, v)
+        average_grads.append(grad_and_var)
+    
+    return average_grads
+    
+
+def get_available_gpus():
+    from tensorflow.python.client import device_lib
+    local_device_protos = device_lib.list_local_devices()
+    return [x.name for x in local_device_protos if x.device_type == 'GPU']
diff --git a/networks/models/train.py b/networks/models/train.py
index 0b52f69..0d37bef 100644
--- a/networks/models/train.py
+++ b/networks/models/train.py
@@ -1,74 +1,79 @@
-import os
-import time
-import numpy as np
-import tensorflow as tf
-import dcgan
-from utils import save_checkpoint, load_checkpoint
-
-def train_dcgan(data, config):
-
-    training_graph = tf.Graph()
-
-    with training_graph.as_default():
-
-        gan = dcgan.dcgan(output_size=config.output_size,
-                          batch_size=config.batch_size,
-                          nd_layers=config.nd_layers,
-                          ng_layers=config.ng_layers,
-                          df_dim=config.df_dim,
-                          gf_dim=config.gf_dim,
-                          c_dim=config.c_dim,
-                          z_dim=config.z_dim,
-                          flip_labels=config.flip_labels,
-                          data_format=config.data_format,
-                          transpose_b=config.transpose_matmul_b)
-
-        save_every_step = True if config.save_every_step == 'True' else False
-
-        gan.training_graph()
-        update_op = gan.optimizer(config.learning_rate, config.beta1)
-
-        checkpoint_dir = os.path.join(config.checkpoint_dir, config.experiment)
-
-        with tf.Session() as sess:
-            writer = tf.summary.FileWriter('./logs/'+config.experiment+'/train', sess.graph)
-
-            init_op = tf.global_variables_initializer()
-            sess.run(init_op)
-
-            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, step=save_every_step)
-
-            epoch = sess.run(gan.increment_epoch)
-            start_time = time.time()
-            for epoch in range(epoch, epoch + config.epoch):
-
-                perm = np.random.permutation(data.shape[0])
-                num_batches = data.shape[0] // config.batch_size
-
-                for idx in range(0, num_batches):
-                    batch_images = data[perm[idx*config.batch_size:(idx+1)*config.batch_size]]
-
-                    _, g_sum, d_sum = sess.run([update_op, gan.g_summary, gan.d_summary], 
-                                               feed_dict={gan.images: batch_images})
-
-                    global_step = gan.global_step.eval()
-                    writer.add_summary(g_sum, global_step)
-                    writer.add_summary(d_sum, global_step)
-
-                    if save_every_step:
-                        save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, global_step, step=True)
-
-                    if config.verbose:
-                        errD_fake = gan.d_loss_fake.eval()
-                        errD_real = gan.d_loss_real.eval({gan.images: batch_images})
-                        errG = gan.g_loss.eval()
-
-                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" \
-                                % (epoch, idx, num_batches, time.time() - start_time, errD_fake+errD_real, errG))
-
-                    elif global_step%100 == 0:
-                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f"%(epoch, idx, num_batches, time.time() - start_time))
-
-                # save a checkpoint every epoch
-                save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, epoch)
-                sess.run(gan.increment_epoch)
+import os
+import sys
+import time
+import numpy as np
+import tensorflow as tf
+import dcgan
+from utils import save_checkpoint, load_checkpoint
+
+
+def train_dcgan(data, config):
+    training_graph = tf.Graph()
+
+    with training_graph.as_default(), tf.device("/cpu:0"):
+        gan = dcgan.DCGAN(output_size=config.output_size,
+                          batch_size=config.batch_size,
+                          nd_layers=config.nd_layers,
+                          ng_layers=config.ng_layers,
+                          df_dim=config.df_dim,
+                          gf_dim=config.gf_dim,
+                          c_dim=config.c_dim,
+                          z_dim=config.z_dim,
+                          flip_labels=config.flip_labels,
+                          data_format=config.data_format,
+                          transpose_b=config.transpose_matmul_b,
+                          num_gpus=config.num_gpus,
+                          learning_rate=config.learning_rate,
+                          beta1=config.beta1)
+
+        save_every_step = True if config.save_every_step == 'True' else False
+
+        g_grads_vars, d_grads_vars = gan.training_graph()
+        update_op = self.apply_gradients(g_grads_vars, d_grads_vars)
+
+        checkpoint_dir = os.path.join(config.checkpoint_dir, config.experiment)
+
+        with tf.Session() as sess:
+            writer = tf.summary.FileWriter('./logs/'+config.experiment+'/train', sess.graph)
+
+            init_op = tf.global_variables_initializer()
+            sess.run(init_op)
+
+            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, step=save_every_step)
+
+            epoch = sess.run(gan.increment_epoch)
+            start_time = time.time()
+            for epoch in range(epoch, epoch + config.epoch):
+
+                perm = np.random.permutation(data.shape[0])
+                num_batches = data.shape[0] // config.batch_size
+
+                for idx in range(0, num_batches):
+                    batch_images = data[perm[idx*config.batch_size:(idx+1)*config.batch_size]]
+
+                    _, g_sum, d_sum = sess.run([update_op, gan.g_summary, gan.d_summary],
+                                               feed_dict={gan.images: batch_images})
+
+                    global_step = gan.global_step.eval()
+                    writer.add_summary(g_sum, global_step)
+                    writer.add_summary(d_sum, global_step)
+
+                    if save_every_step:
+                        save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, global_step, step=True)
+
+                    if config.verbose:
+                        errD_fake = gan.d_loss_fake.eval()
+                        errD_real = gan.d_loss_real.eval({gan.images: batch_images})
+                        errG = gan.g_loss.eval()
+
+                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f"
+                              % (epoch, idx, num_batches, time.time() - start_time, errD_fake + errD_real, errG))
+
+                    elif global_step % 100 == 0:
+                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f" % (epoch, idx, num_batches, time.time() - start_time))
+
+                # save a checkpoint every epoch
+                save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, epoch)
+                sess.run(gan.increment_epoch)
+
+
diff --git a/networks/models/utils.py b/networks/models/utils.py
index ed8e81b..86b80f9 100644
--- a/networks/models/utils.py
+++ b/networks/models/utils.py
@@ -1,41 +1,41 @@
-import os
-import tensorflow as tf
-
-def save_checkpoint(sess, saver, tag, checkpoint_dir, counter, step=False):
-
-    model_name = tag + '.model-'
-    if step:
-        model_name += 'step'
-    else:
-        model_name += 'epoch'
-
-    if not os.path.exists(checkpoint_dir):
-        os.makedirs(checkpoint_dir)
-
-    saver.save(sess, os.path.join(checkpoint_dir, model_name), global_step=counter)
-
-def load_checkpoint(sess, saver, tag, checkpoint_dir, counter=None, step=False):
-    print(" [*] Reading checkpoints...")
-
-    if step:
-        counter_name = 'step'
-    else:
-        counter_name = 'epoch'
-
-    ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
-    if ckpt and ckpt.model_checkpoint_path:
-        ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
-
-        if not counter==None:
-            ckpt_name_epoch = ckpt_name[:ckpt_name.find(counter_name)] + counter_name + '-%i'%counter
-            if os.path.exists(os.path.join(checkpoint_dir, ckpt_name_epoch+'.index')):
-                ckpt_name = ckpt_name_epoch
-            else:
-                print("Checkpoint for ", counter_name , counter_name, "doesn't exist. Using latest checkpoint instead!")
-
-        saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
-        print(" [*] Success to read {}".format(ckpt_name))
-        return True
-    else:
-        print(" [*] Failed to find a checkpoint")
-        return False
+import os
+import tensorflow as tf
+
+def save_checkpoint(sess, saver, tag, checkpoint_dir, counter, step=False):
+
+    model_name = tag + '.model-'
+    if step:
+        model_name += 'step'
+    else:
+        model_name += 'epoch'
+
+    if not os.path.exists(checkpoint_dir):
+        os.makedirs(checkpoint_dir)
+
+    saver.save(sess, os.path.join(checkpoint_dir, model_name), global_step=counter)
+
+def load_checkpoint(sess, saver, tag, checkpoint_dir, counter=None, step=False):
+    print(" [*] Reading checkpoints...")
+
+    if step:
+        counter_name = 'step'
+    else:
+        counter_name = 'epoch'
+
+    ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
+    if ckpt and ckpt.model_checkpoint_path:
+        ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
+
+        if not counter==None:
+            ckpt_name_epoch = ckpt_name[:ckpt_name.find(counter_name)] + counter_name + '-%i'%counter
+            if os.path.exists(os.path.join(checkpoint_dir, ckpt_name_epoch+'.index')):
+                ckpt_name = ckpt_name_epoch
+            else:
+                print("Checkpoint for ", counter_name , counter_name, "doesn't exist. Using latest checkpoint instead!")
+
+        saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
+        print(" [*] Success to read {}".format(ckpt_name))
+        return True
+    else:
+        print(" [*] Failed to find a checkpoint")
+        return False
diff --git a/networks/run_dcgan.py b/networks/run_dcgan.py
index 05895a1..5cbd5c1 100644
--- a/networks/run_dcgan.py
+++ b/networks/run_dcgan.py
@@ -1,37 +1,37 @@
-import os
-import subprocess
-
-datafile = 'data/cosmogan_maps_256_8k_1.npy'
-output_size = 256
-epoch = 50
-flip_labels = 0.01
-batch_size = 64
-z_dim = 64
-nd_layers = 4
-ng_layers = 4
-gf_dim = 64
-df_dim = 64
-save_every_step = 'False'
-data_format = 'NCHW'
-transpose_matmul_b = False
-verbose = 'True'
-
-experiment = 'cosmo_myExp_batchSize%i_flipLabel%0.3f_'\
-             'nd%i_ng%i_gfdim%i_dfdim%i_zdim%i'%(batch_size, flip_labels, nd_layers,\
-                                                 ng_layers, gf_dim, df_dim, z_dim)
-
-command = 'python -m models.main --dataset cosmo --datafile %s '\
-          '--output_size %i --flip_labels %f --experiment %s '\
-          '--epoch %i --batch_size %i --z_dim %i '\
-          '--nd_layers %i --ng_layers %i --gf_dim %i --df_dim %i --save_every_step %s '\
-          '--data_format %s --transpose_matmul_b %s --verbose %s'%(datafile, output_size, flip_labels, experiment,\
-                                                                   epoch, batch_size, z_dim,\
-                                                                   nd_layers, ng_layers, gf_dim, df_dim, save_every_step,\
-                                                                   data_format, transpose_matmul_b, verbose)
-
-if not os.path.isdir('output'):
-    os.mkdir('output')
-
-print command.split()
-f_out = open('output/'+experiment+'.log', 'w')
-subprocess.call(command.split(), stdout=f_out)
+import os
+import subprocess
+
+datafile = 'data/cosmogan_maps_256_8k_1.npy'
+output_size = 256
+epoch = 50
+flip_labels = 0.01
+batch_size = 64
+z_dim = 64
+nd_layers = 4
+ng_layers = 4
+gf_dim = 64
+df_dim = 64
+save_every_step = 'False'
+data_format = 'NCHW'
+transpose_matmul_b = False
+verbose = 'True'
+
+experiment = 'cosmo_myExp_batchSize%i_flipLabel%0.3f_'\
+             'nd%i_ng%i_gfdim%i_dfdim%i_zdim%i'%(batch_size, flip_labels, nd_layers,\
+                                                 ng_layers, gf_dim, df_dim, z_dim)
+
+command = 'python -m models.main --dataset cosmo --datafile %s '\
+          '--output_size %i --flip_labels %f --experiment %s '\
+          '--epoch %i --batch_size %i --z_dim %i '\
+          '--nd_layers %i --ng_layers %i --gf_dim %i --df_dim %i --save_every_step %s '\
+          '--data_format %s --transpose_matmul_b %s --verbose %s'%(datafile, output_size, flip_labels, experiment,\
+                                                                   epoch, batch_size, z_dim,\
+                                                                   nd_layers, ng_layers, gf_dim, df_dim, save_every_step,\
+                                                                   data_format, transpose_matmul_b, verbose)
+
+if not os.path.isdir('output'):
+    os.mkdir('output')
+
+print command.split()
+f_out = open('output/'+experiment+'.log', 'w')
+subprocess.call(command.split(), stdout=f_out)

From cd0ed4af314d679b0ee4773932feee48ed4499a6 Mon Sep 17 00:00:00 2001
From: Wesley Emeneker <wesley.d.emeneker.ctr@mhpcc.hpc.mil>
Date: Mon, 27 Aug 2018 12:15:59 -1000
Subject: [PATCH 3/7] Reworked cmdline argument parsing and setting defaults.

Now it is easy to change the defaults.
---
 networks/models/dcgan.py | 25 +++++++++++-----
 networks/models/main.py  | 16 ++++++++--
 networks/models/ops.py   | 18 +++++++++---
 networks/models/train.py |  8 ++---
 networks/run_dcgan.py    | 63 ++++++++++++++++++++--------------------
 5 files changed, 82 insertions(+), 48 deletions(-)

diff --git a/networks/models/dcgan.py b/networks/models/dcgan.py
index 302b612..8f6afb5 100644
--- a/networks/models/dcgan.py
+++ b/networks/models/dcgan.py
@@ -77,7 +77,6 @@ def training_graph(self):
                         with tf.variable_scope("generator"):
                             g_images = self.generator(z_images, is_training=True)
 
-                        print >>sys.stderr, "SHAPES:", g_images.shape, images.shape, z_images.shape, self.batch_size, self.compute_devices
                         with tf.variable_scope("discriminator") as d_scope:
                             d_scope.reuse_variables()
                             d_prob_fake, d_logits_fake = self.discriminator(g_images, is_training=True)
@@ -92,9 +91,15 @@ def training_graph(self):
                             with tf.name_scope("g"):
                                 g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_logits_fake)))
 
+                        # Each tower should reuse the same variables
+                        #   across towers so that when we apply the 
+                        #   averaged gradients, we only have to update each variable
+                        #   once, instead of each variable per tower.
+                        tf.get_variable_scope().reuse_variables()
+                        
                         t_vars = tf.trainable_variables()
-                        d_vars = [var for var in t_vars if 'discriminator/' in var.name and "gradients" in var.name]
-                        g_vars = [var for var in t_vars if 'generator/' in var.name and "gradients" in var.names]
+                        d_vars = [var for var in t_vars if 'discriminator/' in var.name]
+                        g_vars = [var for var in t_vars if 'generator/' in var.name]
 
                         d_gradients, g_gradients = self.compute_gradients(d_loss, g_loss, d_vars, g_vars)
                         tower_d_gradients.append(d_gradients)
@@ -105,7 +110,15 @@ def training_graph(self):
         # I think so.
         g_grads_vars = average_gradients(tower_g_gradients)
         d_grads_vars = average_gradients(tower_d_gradients)
-                        
+
+        # FIXME: These variables should be removed once I figure out how to make them
+        #   available through the summary.
+        # FIXME: These variables only hold the last result from the data-parallel GPU
+        #   operations. 
+        self.d_loss_fake = d_loss_fake
+        self.d_loss_real = d_loss_real
+        self.g_loss = g_loss
+        
         self.d_summary = tf.summary.merge([tf.summary.histogram("prob/real", d_prob_real),
                                            tf.summary.histogram("prob/fake", d_prob_fake),
                                            tf.summary.scalar("loss/real", d_loss_real),
@@ -158,7 +171,7 @@ def compute_gradients(self, d_loss, g_loss, d_vars, g_vars):
 
         return d_optim, g_optim
 
-    def apply_gradients(g_grads_vars, d_grads_vars):
+    def apply_gradients(self, g_grads_vars, d_grads_vars):
         return tf.group(self.g_optim.apply_gradients(g_grads_vars),
                         self.d_optim.apply_gradients(d_grads_vars, global_step=self.global_step))
 
@@ -173,8 +186,6 @@ def generator(self, z, is_training):
 
         chain = h0
 
-        print >>sys.stderr, "chain shape:", chain.shape
-        
         for h in range(1, self.ng_layers):
             map_size *= self.stride
             num_channels /= 2
diff --git a/networks/models/main.py b/networks/models/main.py
index af9c676..54d8357 100644
--- a/networks/models/main.py
+++ b/networks/models/main.py
@@ -28,8 +28,20 @@
 config = flags.FLAGS
 
 
-def main(_):
-    pprint.PrettyPrinter().pprint(config.__flags)
+def main(**config_kwargs):
+    import sys
+    # config_kwargs are the default arguments for everything.
+    # These can be overridden by anything specified on the command-line
+    for k, v in config_kwargs.iteritems():
+        try:
+            config[k].value = v
+        except flags.UnrecognizedFlagError:
+            pass
+
+    # Now override the defaults from anything specified on the command-line
+    args = flags.FLAGS(sys.argv, known_only=False)
+    
+    pprint.PrettyPrinter().pprint({k: config[k].value for k in config.__flags })
     train.train_dcgan(get_data(), config)
 
 
diff --git a/networks/models/ops.py b/networks/models/ops.py
index 2d1ddc5..30e9919 100644
--- a/networks/models/ops.py
+++ b/networks/models/ops.py
@@ -1,3 +1,5 @@
+from collections import defaultdict
+import itertools
 from pprint import pprint
 import sys
 
@@ -96,16 +98,23 @@ def average_gradients(tower_grads):
      List of pairs of (gradient, variable) where the gradient has been averaged
      across all towers.
     """
-    pprint(tower_grads, sys.stderr)
+    # The variables we get might not have a match from each tower
+    #   or the variable order from each tower may not be the same.
+    # So, we need to do our own grouping here.
+    grouped_vars = defaultdict(list)
+    for grad_and_var in itertools.chain(*tower_grads):
+        # Remove just the "tower_N/" part of the variable name.
+        grouped_var_name = '/'.join(grad_and_var[1].name.split("/")[1:])
+        grouped_vars[grouped_var_name].append(grad_and_var)
+    
     average_grads = []
-    for grad_and_vars in zip(*tower_grads):
+    for grad_and_vars in grouped_vars.itervalues():
         # Note that each grad_and_vars looks like the following:
         #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
         grads = []
         for g, _ in grad_and_vars:
-            print >>sys.stderr, "var name:", _.name, _
             if g is None:
-                pass
+                continue
             # Add 0 dimension to the gradients to represent the tower.
             expanded_g = tf.expand_dims(g, 0)
 
@@ -119,6 +128,7 @@ def average_gradients(tower_grads):
         # Keep in mind that the Variables are redundant because they are shared
         # across towers. So .. we will just return the first tower's pointer to
         # the Variable.
+        # FIXME: This may not be true for the cosmoGAN towers
         v = grad_and_vars[0][1]
         grad_and_var = (grad, v)
         average_grads.append(grad_and_var)
diff --git a/networks/models/train.py b/networks/models/train.py
index 0d37bef..96e85b8 100644
--- a/networks/models/train.py
+++ b/networks/models/train.py
@@ -29,7 +29,7 @@ def train_dcgan(data, config):
         save_every_step = True if config.save_every_step == 'True' else False
 
         g_grads_vars, d_grads_vars = gan.training_graph()
-        update_op = self.apply_gradients(g_grads_vars, d_grads_vars)
+        update_op = gan.apply_gradients(g_grads_vars, d_grads_vars)
 
         checkpoint_dir = os.path.join(config.checkpoint_dir, config.experiment)
 
@@ -66,11 +66,11 @@ def train_dcgan(data, config):
                         errD_real = gan.d_loss_real.eval({gan.images: batch_images})
                         errG = gan.g_loss.eval()
 
-                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f"
-                              % (epoch, idx, num_batches, time.time() - start_time, errD_fake + errD_real, errG))
+                        print ("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f"
+                               % (epoch, idx, num_batches, time.time() - start_time, errD_fake + errD_real, errG))
 
                     elif global_step % 100 == 0:
-                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f" % (epoch, idx, num_batches, time.time() - start_time))
+                        print "Epoch: [%2d] Step: [%4d/%4d] time: %4.4f" % (epoch, idx, num_batches, time.time() - start_time)
 
                 # save a checkpoint every epoch
                 save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, epoch)
diff --git a/networks/run_dcgan.py b/networks/run_dcgan.py
index 5cbd5c1..476345b 100644
--- a/networks/run_dcgan.py
+++ b/networks/run_dcgan.py
@@ -1,37 +1,38 @@
 import os
-import subprocess
 
-datafile = 'data/cosmogan_maps_256_8k_1.npy'
-output_size = 256
-epoch = 50
-flip_labels = 0.01
-batch_size = 64
-z_dim = 64
-nd_layers = 4
-ng_layers = 4
-gf_dim = 64
-df_dim = 64
-save_every_step = 'False'
-data_format = 'NCHW'
-transpose_matmul_b = False
-verbose = 'True'
+# These will be the default parameters for the DCGAN
+config_kwargs = dict(
+    datafile = 'data/cosmogan_maps_256_8k_1.npy',
+    output_size = 256,
+    epoch = 50,
+    flip_labels = 0.01,
+    batch_size = 128,
+    z_dim = 64,
+    nd_layers = 4,
+    ng_layers = 4,
+    gf_dim = 64,
+    df_dim = 64,
+    save_every_step = 'False',
+    data_format = 'NCHW',
+    transpose_matmul_b = False,
+    verbose = 'True',
+)
 
-experiment = 'cosmo_myExp_batchSize%i_flipLabel%0.3f_'\
-             'nd%i_ng%i_gfdim%i_dfdim%i_zdim%i'%(batch_size, flip_labels, nd_layers,\
-                                                 ng_layers, gf_dim, df_dim, z_dim)
+config_kwargs['experiment'] = ('cosmo_myExp_batchSize%i_flipLabel%0.3f_nd%i_ng%i_gfdim%i_dfdim%i_zdim%i' %
+                               (config_kwargs['batch_size'],
+                                config_kwargs['flip_labels'],
+                                config_kwargs['nd_layers'],
+                                config_kwargs['ng_layers'],
+                                config_kwargs['gf_dim'],
+                                config_kwargs['df_dim'],
+                                config_kwargs['z_dim']))
 
-command = 'python -m models.main --dataset cosmo --datafile %s '\
-          '--output_size %i --flip_labels %f --experiment %s '\
-          '--epoch %i --batch_size %i --z_dim %i '\
-          '--nd_layers %i --ng_layers %i --gf_dim %i --df_dim %i --save_every_step %s '\
-          '--data_format %s --transpose_matmul_b %s --verbose %s'%(datafile, output_size, flip_labels, experiment,\
-                                                                   epoch, batch_size, z_dim,\
-                                                                   nd_layers, ng_layers, gf_dim, df_dim, save_every_step,\
-                                                                   data_format, transpose_matmul_b, verbose)
 
-if not os.path.isdir('output'):
-    os.mkdir('output')
+if __name__ == "__main__":
+    import models.main
+    import sys
+    
+    if not os.path.isdir('output'):
+        os.mkdir('output')
 
-print command.split()
-f_out = open('output/'+experiment+'.log', 'w')
-subprocess.call(command.split(), stdout=f_out)
+    models.main.main(**config_kwargs)

From e8451e85fd874fe05cdac8a5ee38f9e78802e48b Mon Sep 17 00:00:00 2001
From: Wesley Emeneker <wesley.d.emeneker.ctr@mhpcc.hpc.mil>
Date: Mon, 27 Aug 2018 12:23:58 -1000
Subject: [PATCH 4/7] .pyc files ignored by git.

---
 .gitignore | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 0d20b64..94487b9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1 @@
-*.pyc
+*.pyc

From 7e67db4705aacd0fab88b6e5bb52a95ea0da2b3a Mon Sep 17 00:00:00 2001
From: Wesley Emeneker <wesley.d.emeneker.ctr@mhpcc.hpc.mil>
Date: Tue, 28 Aug 2018 13:32:41 -1000
Subject: [PATCH 5/7] Added option to write output to stdout, or to an output
 file.

---
 networks/run_dcgan.py | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/networks/run_dcgan.py b/networks/run_dcgan.py
index 476345b..288b858 100644
--- a/networks/run_dcgan.py
+++ b/networks/run_dcgan.py
@@ -31,8 +31,25 @@
 if __name__ == "__main__":
     import models.main
     import sys
-    
-    if not os.path.isdir('output'):
-        os.mkdir('output')
 
-    models.main.main(**config_kwargs)
+    # The default is to log stdout to an output file.
+    # Only if the user gives the "--stdout" command line option
+    #   should we display stdout to the real stdout.
+    _old_stdout = sys.stdout
+    _new_stdout = sys.stdout
+
+    if "--stdout" in sys.argv:
+        sys.argv.remove("--stdout")
+
+    else:        
+        if not os.path.isdir('output'):
+            os.mkdir('output')
+            
+        _new_stdout = open(os.path.join("output", config_kwargs['experiment']), 'wb')
+        
+
+    try:
+        sys.stdout = _new_stdout
+        models.main.main(**config_kwargs)
+    finally:
+        sys.stdout = _old_stdout

From 6990e807b702b62bbdc6528da0d825654a0a2b7c Mon Sep 17 00:00:00 2001
From: Wesley Emeneker <wesley.d.emeneker.ctr@mhpcc.hpc.mil>
Date: Mon, 24 Sep 2018 14:43:10 -1000
Subject: [PATCH 6/7] Windows->Unix line endings

---
 networks/models/dcgan.py | 508 +++++++++++++++++++--------------------
 1 file changed, 254 insertions(+), 254 deletions(-)

diff --git a/networks/models/dcgan.py b/networks/models/dcgan.py
index 8f6afb5..df904d5 100644
--- a/networks/models/dcgan.py
+++ b/networks/models/dcgan.py
@@ -1,254 +1,254 @@
-import sys
-
-import tensorflow as tf
-from .ops import linear, conv2d, conv2d_transpose, lrelu, average_gradients, get_available_gpus
-
-
-class DCGAN(object):
-    def __init__(self, output_size=64, batch_size=64,
-                 nd_layers=4, ng_layers=4, df_dim=128, gf_dim=128,
-                 c_dim=1, z_dim=100, flip_labels=0.01, data_format="NHWC",
-                 gen_prior=tf.random_normal, transpose_b=False,
-                 num_gpus=-1, learning_rate=.0002, beta1=0.5):
-
-        self.output_size = output_size
-        self.batch_size = batch_size
-        self.nd_layers = nd_layers # Number of hidden layers in the discriminator?
-        self.ng_layers = ng_layers # Number of hidden layers in the generator?
-        self.df_dim = df_dim # discriminator f dimensions?
-        self.gf_dim = gf_dim # generator f dimensions
-        self.c_dim = c_dim # number of channels
-        self.z_dim = z_dim # 
-        self.flip_labels = flip_labels
-        self.data_format = data_format
-        self.gen_prior = gen_prior
-        self.transpose_b = transpose_b  # transpose weight matrix in linear layers for (possible) better performance when running on HSW/KNL
-        self.stride = 2  # this is fixed for this architecture
-        
-        self.compute_devices = get_available_gpus()
-        # Figure out if we are using gpus, how many, which ones
-        # If num_gpus < 0, use all GPUs we were given
-        if num_gpus > 0:
-            # If the user doesn't want to use all the gpus, take the first N
-            self.compute_devices = self.compute_devices[:num_gpus]
-        elif num_gpus == 0:
-            self.compute_devices = ['/cpu:0']
-
-        self._check_architecture_consistency()
-
-        self.batchnorm_kwargs = {'epsilon': 1e-5, 'decay': 0.9,
-                                 'updates_collections': None, 'scale': True,
-                                 'fused': True, 'data_format': self.data_format}
-        self.make_optimizers(learning_rate, beta1)
-
-    @property
-    def compute_batch_size(self):
-        """Return the batch size that each GPU (or other compute device) should have"""
-        return self.batch_size / len(self.compute_devices)
-        
-    def training_graph(self):
-        """Make the training graph such that it divides the work among 1 or more GPUS"""
-
-        if self.data_format == "NHWC":
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
-        else:
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
-            
-        self.z = self.gen_prior(shape=[self.batch_size, self.z_dim])
-
-        # Take the batch of images and split the generator and discriminator work among
-        #   all the computing devices given.
-        # My hope is that by parallelizing this inner graph, we'll get some
-        #   nice speedup without having to refactor a lot.
-        real_image_splits = tf.split(self.images, num_or_size_splits=len(self.compute_devices), axis=0)
-        # Split the random-noise images into chunks, one for each compute device
-        gen_image_splits = tf.split(self.z, num_or_size_splits=len(self.compute_devices), axis=0)
-
-        tower_d_gradients = []
-        tower_g_gradients = []
-        with tf.variable_scope(tf.get_variable_scope()):
-            for idx, (images, z_images) in enumerate(zip(real_image_splits, gen_image_splits)):
-                compute_device = self.compute_devices[idx]
-                with tf.device(compute_device):
-                    with tf.variable_scope("tower_%d" % idx) as scope:
-                        with tf.variable_scope("discriminator"):
-                            d_prob_real, d_logits_real = self.discriminator(images, is_training=True)
-
-                        with tf.variable_scope("generator"):
-                            g_images = self.generator(z_images, is_training=True)
-
-                        with tf.variable_scope("discriminator") as d_scope:
-                            d_scope.reuse_variables()
-                            d_prob_fake, d_logits_fake = self.discriminator(g_images, is_training=True)
-
-                        with tf.name_scope("losses"):
-                            with tf.name_scope("d"):
-                                d_label_real, d_label_fake = self._labels()
-                                d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real, labels=d_label_real, name="real"))
-                                d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=d_label_fake, name="fake"))
-                                d_loss = d_loss_real + d_loss_fake
-
-                            with tf.name_scope("g"):
-                                g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_logits_fake)))
-
-                        # Each tower should reuse the same variables
-                        #   across towers so that when we apply the 
-                        #   averaged gradients, we only have to update each variable
-                        #   once, instead of each variable per tower.
-                        tf.get_variable_scope().reuse_variables()
-                        
-                        t_vars = tf.trainable_variables()
-                        d_vars = [var for var in t_vars if 'discriminator/' in var.name]
-                        g_vars = [var for var in t_vars if 'generator/' in var.name]
-
-                        d_gradients, g_gradients = self.compute_gradients(d_loss, g_loss, d_vars, g_vars)
-                        tower_d_gradients.append(d_gradients)
-                        tower_g_gradients.append(g_gradients)
-
-        # Do these variables need to be separated by variable name so that
-        #   the average_gradients function will correctly aggregate them?
-        # I think so.
-        g_grads_vars = average_gradients(tower_g_gradients)
-        d_grads_vars = average_gradients(tower_d_gradients)
-
-        # FIXME: These variables should be removed once I figure out how to make them
-        #   available through the summary.
-        # FIXME: These variables only hold the last result from the data-parallel GPU
-        #   operations. 
-        self.d_loss_fake = d_loss_fake
-        self.d_loss_real = d_loss_real
-        self.g_loss = g_loss
-        
-        self.d_summary = tf.summary.merge([tf.summary.histogram("prob/real", d_prob_real),
-                                           tf.summary.histogram("prob/fake", d_prob_fake),
-                                           tf.summary.scalar("loss/real", d_loss_real),
-                                           tf.summary.scalar("loss/fake", d_loss_fake),
-                                           tf.summary.scalar("loss/d", d_loss)])
-
-        g_sum = [tf.summary.scalar("loss/g", g_loss)]
-        if self.data_format == "NHWC":  # tf.summary.image is not implemented for NCHW
-            g_sum.append(tf.summary.image("G", g_images, max_outputs=4))
-        self.g_summary = tf.summary.merge(g_sum)
-
-        with tf.variable_scope("counters"):
-            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
-            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
-            self.global_step = tf.train.get_or_create_global_step()
-
-        self.saver = tf.train.Saver(max_to_keep=8000)
-
-        return g_grads_vars, d_grads_vars
-        
-    def inference_graph(self):
-
-        if self.data_format == "NHWC":
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
-        else:
-            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
-
-        self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
-
-        with tf.variable_scope("discriminator"):
-            self.D, _ = self.discriminator(self.images, is_training=False)
-
-        with tf.variable_scope("generator"):
-            self.G = self.generator(self.z, is_training=False)
-
-        with tf.variable_scope("counters"):
-            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
-            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
-            self.global_step = tf.train.get_or_create_global_step()
-
-        self.saver = tf.train.Saver(max_to_keep=8000)
-
-    def make_optimizers(self, learning_rate, beta1):
-        self.d_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1)
-        self.g_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1)
-        
-    def compute_gradients(self, d_loss, g_loss, d_vars, g_vars):
-        d_optim = self.d_optim.compute_gradients(d_loss, var_list=d_vars)
-        g_optim = self.g_optim.compute_gradients(g_loss, var_list=g_vars)
-
-        return d_optim, g_optim
-
-    def apply_gradients(self, g_grads_vars, d_grads_vars):
-        return tf.group(self.g_optim.apply_gradients(g_grads_vars),
-                        self.d_optim.apply_gradients(d_grads_vars, global_step=self.global_step))
-
-    def generator(self, z, is_training):
-        map_size = self.output_size / int(2**self.ng_layers)
-        num_channels = self.gf_dim * int(2**(self.ng_layers - 1))
-
-        z_ = linear(z, num_channels*map_size*map_size, 'h0_lin', transpose_b=self.transpose_b)
-        h0 = tf.reshape(z_, self._tensor_data_format(-1, map_size, map_size, num_channels))
-        bn0 = tf.contrib.layers.batch_norm(h0, is_training=is_training, scope='bn0', **self.batchnorm_kwargs)
-        h0 = tf.nn.relu(bn0)
-
-        chain = h0
-
-        for h in range(1, self.ng_layers):
-            map_size *= self.stride
-            num_channels /= 2
-            chain = conv2d_transpose(chain,
-                                     self._tensor_data_format(self.compute_batch_size, map_size, map_size, num_channels),
-                                     stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % h)
-            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
-            chain = tf.nn.relu(chain)
-
-        map_size *= self.stride
-        hn = conv2d_transpose(chain,
-                              self._tensor_data_format(self.compute_batch_size, map_size, map_size, self.c_dim),
-                              stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % self.ng_layers)
-
-        return tf.nn.tanh(hn)
-
-    def discriminator(self, image, is_training):
-        """This discriminator works on a single image at a time"""
-
-        chain = lrelu(conv2d(image, self.df_dim, self.data_format, name='h0_conv'))
-
-        for h in range(1, self.nd_layers):
-            chain = conv2d(chain, self.df_dim*(2**h), self.data_format, name='h%i_conv' % h)
-            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
-            chain = lrelu(chain)
-
-        hn = linear(tf.reshape(chain, [self.batch_size, -1]),
-                    1,
-                    'h%i_lin' % self.nd_layers,
-                    transpose_b=self.transpose_b)
-
-        return tf.nn.sigmoid(hn), hn
-
-    def _labels(self):
-        with tf.name_scope("labels"):
-            ones = tf.ones([self.batch_size, 1])
-            zeros = tf.zeros([self.batch_size, 1])
-            flip_labels = tf.constant(self.flip_labels)
-
-            if self.flip_labels > 0:
-                prob = tf.random_uniform([], 0, 1)
-
-                d_label_real = tf.cond(tf.less(prob, flip_labels), lambda: zeros, lambda: ones)
-                d_label_fake = tf.cond(tf.less(prob, flip_labels), lambda: ones, lambda: zeros)
-            else:
-                d_label_real = ones
-                d_label_fake = zeros
-
-        return d_label_real, d_label_fake
-
-    def _tensor_data_format(self, N, H, W, C):
-        if self.data_format == "NHWC":
-            return [int(N), int(H), int(W), int(C)]
-        else:
-            return [int(N), int(C), int(H), int(W)]
-
-    def _check_architecture_consistency(self):
-
-        if self.output_size/2**self.nd_layers < 1:
-            print("Error: Number of discriminator conv. layers are larger than the output_size for this architecture")
-            exit(0)
-
-        if self.output_size/2**self.ng_layers < 1:
-            print("Error: Number of generator conv_transpose layers are larger than the output_size for this architecture")
-            exit(0)
-
+import sys
+
+import tensorflow as tf
+from .ops import linear, conv2d, conv2d_transpose, lrelu, average_gradients, get_available_gpus
+
+
+class DCGAN(object):
+    def __init__(self, output_size=64, batch_size=64,
+                 nd_layers=4, ng_layers=4, df_dim=128, gf_dim=128,
+                 c_dim=1, z_dim=100, flip_labels=0.01, data_format="NHWC",
+                 gen_prior=tf.random_normal, transpose_b=False,
+                 num_gpus=-1, learning_rate=.0002, beta1=0.5):
+
+        self.output_size = output_size
+        self.batch_size = batch_size
+        self.nd_layers = nd_layers # Number of hidden layers in the discriminator?
+        self.ng_layers = ng_layers # Number of hidden layers in the generator?
+        self.df_dim = df_dim # discriminator f dimensions?
+        self.gf_dim = gf_dim # generator f dimensions
+        self.c_dim = c_dim # number of channels
+        self.z_dim = z_dim # 
+        self.flip_labels = flip_labels
+        self.data_format = data_format
+        self.gen_prior = gen_prior
+        self.transpose_b = transpose_b  # transpose weight matrix in linear layers for (possible) better performance when running on HSW/KNL
+        self.stride = 2  # this is fixed for this architecture
+        
+        self.compute_devices = get_available_gpus()
+        # Figure out if we are using gpus, how many, which ones
+        # If num_gpus < 0, use all GPUs we were given
+        if num_gpus > 0:
+            # If the user doesn't want to use all the gpus, take the first N
+            self.compute_devices = self.compute_devices[:num_gpus]
+        elif num_gpus == 0:
+            self.compute_devices = ['/cpu:0']
+
+        self._check_architecture_consistency()
+
+        self.batchnorm_kwargs = {'epsilon': 1e-5, 'decay': 0.9,
+                                 'updates_collections': None, 'scale': True,
+                                 'fused': True, 'data_format': self.data_format}
+        self.make_optimizers(learning_rate, beta1)
+
+    @property
+    def compute_batch_size(self):
+        """Return the batch size that each GPU (or other compute device) should have"""
+        return self.batch_size / len(self.compute_devices)
+        
+    def training_graph(self):
+        """Make the training graph such that it divides the work among 1 or more GPUS"""
+
+        if self.data_format == "NHWC":
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
+        else:
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
+            
+        self.z = self.gen_prior(shape=[self.batch_size, self.z_dim])
+
+        # Take the batch of images and split the generator and discriminator work among
+        #   all the computing devices given.
+        # My hope is that by parallelizing this inner graph, we'll get some
+        #   nice speedup without having to refactor a lot.
+        real_image_splits = tf.split(self.images, num_or_size_splits=len(self.compute_devices), axis=0)
+        # Split the random-noise images into chunks, one for each compute device
+        gen_image_splits = tf.split(self.z, num_or_size_splits=len(self.compute_devices), axis=0)
+
+        tower_d_gradients = []
+        tower_g_gradients = []
+        with tf.variable_scope(tf.get_variable_scope()):
+            for idx, (images, z_images) in enumerate(zip(real_image_splits, gen_image_splits)):
+                compute_device = self.compute_devices[idx]
+                with tf.device(compute_device):
+                    with tf.variable_scope("tower_%d" % idx) as scope:
+                        with tf.variable_scope("discriminator"):
+                            d_prob_real, d_logits_real = self.discriminator(images, is_training=True)
+
+                        with tf.variable_scope("generator"):
+                            g_images = self.generator(z_images, is_training=True)
+
+                        with tf.variable_scope("discriminator") as d_scope:
+                            d_scope.reuse_variables()
+                            d_prob_fake, d_logits_fake = self.discriminator(g_images, is_training=True)
+
+                        with tf.name_scope("losses"):
+                            with tf.name_scope("d"):
+                                d_label_real, d_label_fake = self._labels()
+                                d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real, labels=d_label_real, name="real"))
+                                d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=d_label_fake, name="fake"))
+                                d_loss = d_loss_real + d_loss_fake
+
+                            with tf.name_scope("g"):
+                                g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_logits_fake)))
+
+                        # Each tower should reuse the same variables
+                        #   across towers so that when we apply the 
+                        #   averaged gradients, we only have to update each variable
+                        #   once, instead of each variable per tower.
+                        tf.get_variable_scope().reuse_variables()
+                        
+                        t_vars = tf.trainable_variables()
+                        d_vars = [var for var in t_vars if 'discriminator/' in var.name]
+                        g_vars = [var for var in t_vars if 'generator/' in var.name]
+
+                        d_gradients, g_gradients = self.compute_gradients(d_loss, g_loss, d_vars, g_vars)
+                        tower_d_gradients.append(d_gradients)
+                        tower_g_gradients.append(g_gradients)
+
+        # Do these variables need to be separated by variable name so that
+        #   the average_gradients function will correctly aggregate them?
+        # I think so.
+        g_grads_vars = average_gradients(tower_g_gradients)
+        d_grads_vars = average_gradients(tower_d_gradients)
+
+        # FIXME: These variables should be removed once I figure out how to make them
+        #   available through the summary.
+        # FIXME: These variables only hold the last result from the data-parallel GPU
+        #   operations. 
+        self.d_loss_fake = d_loss_fake
+        self.d_loss_real = d_loss_real
+        self.g_loss = g_loss
+        
+        self.d_summary = tf.summary.merge([tf.summary.histogram("prob/real", d_prob_real),
+                                           tf.summary.histogram("prob/fake", d_prob_fake),
+                                           tf.summary.scalar("loss/real", d_loss_real),
+                                           tf.summary.scalar("loss/fake", d_loss_fake),
+                                           tf.summary.scalar("loss/d", d_loss)])
+
+        g_sum = [tf.summary.scalar("loss/g", g_loss)]
+        if self.data_format == "NHWC":  # tf.summary.image is not implemented for NCHW
+            g_sum.append(tf.summary.image("G", g_images, max_outputs=4))
+        self.g_summary = tf.summary.merge(g_sum)
+
+        with tf.variable_scope("counters"):
+            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
+            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
+            self.global_step = tf.train.get_or_create_global_step()
+
+        self.saver = tf.train.Saver(max_to_keep=8000)
+
+        return g_grads_vars, d_grads_vars
+        
+    def inference_graph(self):
+
+        if self.data_format == "NHWC":
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
+        else:
+            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')
+
+        self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
+
+        with tf.variable_scope("discriminator"):
+            self.D, _ = self.discriminator(self.images, is_training=False)
+
+        with tf.variable_scope("generator"):
+            self.G = self.generator(self.z, is_training=False)
+
+        with tf.variable_scope("counters"):
+            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
+            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
+            self.global_step = tf.train.get_or_create_global_step()
+
+        self.saver = tf.train.Saver(max_to_keep=8000)
+
+    def make_optimizers(self, learning_rate, beta1):
+        self.d_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1)
+        self.g_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1)
+        
+    def compute_gradients(self, d_loss, g_loss, d_vars, g_vars):
+        d_optim = self.d_optim.compute_gradients(d_loss, var_list=d_vars)
+        g_optim = self.g_optim.compute_gradients(g_loss, var_list=g_vars)
+
+        return d_optim, g_optim
+
+    def apply_gradients(self, g_grads_vars, d_grads_vars):
+        return tf.group(self.g_optim.apply_gradients(g_grads_vars),
+                        self.d_optim.apply_gradients(d_grads_vars, global_step=self.global_step))
+
+    def generator(self, z, is_training):
+        map_size = self.output_size / int(2**self.ng_layers)
+        num_channels = self.gf_dim * int(2**(self.ng_layers - 1))
+
+        z_ = linear(z, num_channels*map_size*map_size, 'h0_lin', transpose_b=self.transpose_b)
+        h0 = tf.reshape(z_, self._tensor_data_format(-1, map_size, map_size, num_channels))
+        bn0 = tf.contrib.layers.batch_norm(h0, is_training=is_training, scope='bn0', **self.batchnorm_kwargs)
+        h0 = tf.nn.relu(bn0)
+
+        chain = h0
+
+        for h in range(1, self.ng_layers):
+            map_size *= self.stride
+            num_channels /= 2
+            chain = conv2d_transpose(chain,
+                                     self._tensor_data_format(self.compute_batch_size, map_size, map_size, num_channels),
+                                     stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % h)
+            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
+            chain = tf.nn.relu(chain)
+
+        map_size *= self.stride
+        hn = conv2d_transpose(chain,
+                              self._tensor_data_format(self.compute_batch_size, map_size, map_size, self.c_dim),
+                              stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T' % self.ng_layers)
+
+        return tf.nn.tanh(hn)
+
+    def discriminator(self, image, is_training):
+        """This discriminator works on a single image at a time"""
+
+        chain = lrelu(conv2d(image, self.df_dim, self.data_format, name='h0_conv'))
+
+        for h in range(1, self.nd_layers):
+            chain = conv2d(chain, self.df_dim*(2**h), self.data_format, name='h%i_conv' % h)
+            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i' % h, **self.batchnorm_kwargs)
+            chain = lrelu(chain)
+
+        hn = linear(tf.reshape(chain, [self.batch_size, -1]),
+                    1,
+                    'h%i_lin' % self.nd_layers,
+                    transpose_b=self.transpose_b)
+
+        return tf.nn.sigmoid(hn), hn
+
+    def _labels(self):
+        with tf.name_scope("labels"):
+            ones = tf.ones([self.batch_size, 1])
+            zeros = tf.zeros([self.batch_size, 1])
+            flip_labels = tf.constant(self.flip_labels)
+
+            if self.flip_labels > 0:
+                prob = tf.random_uniform([], 0, 1)
+
+                d_label_real = tf.cond(tf.less(prob, flip_labels), lambda: zeros, lambda: ones)
+                d_label_fake = tf.cond(tf.less(prob, flip_labels), lambda: ones, lambda: zeros)
+            else:
+                d_label_real = ones
+                d_label_fake = zeros
+
+        return d_label_real, d_label_fake
+
+    def _tensor_data_format(self, N, H, W, C):
+        if self.data_format == "NHWC":
+            return [int(N), int(H), int(W), int(C)]
+        else:
+            return [int(N), int(C), int(H), int(W)]
+
+    def _check_architecture_consistency(self):
+
+        if self.output_size/2**self.nd_layers < 1:
+            print("Error: Number of discriminator conv. layers are larger than the output_size for this architecture")
+            exit(0)
+
+        if self.output_size/2**self.ng_layers < 1:
+            print("Error: Number of generator conv_transpose layers are larger than the output_size for this architecture")
+            exit(0)
+

From 5a51132547a703e1b1a08f3b4e3ba00d697e8f7c Mon Sep 17 00:00:00 2001
From: Wesley Emeneker <wesley.d.emeneker.ctr@mhpcc.hpc.mil>
Date: Mon, 24 Sep 2018 14:44:34 -1000
Subject: [PATCH 7/7] Windows->Unix line endings.

---
 networks/models/main.py  | 120 ++++++++---------
 networks/models/ops.py   | 284 +++++++++++++++++++--------------------
 networks/models/train.py | 158 +++++++++++-----------
 networks/models/utils.py |  82 +++++------
 networks/run_dcgan.py    | 110 +++++++--------
 5 files changed, 377 insertions(+), 377 deletions(-)

diff --git a/networks/models/main.py b/networks/models/main.py
index 54d8357..41f8d93 100644
--- a/networks/models/main.py
+++ b/networks/models/main.py
@@ -1,60 +1,60 @@
-import tensorflow as tf
-import train
-import numpy as np
-import pprint
-
-flags = tf.app.flags
-flags.DEFINE_string("dataset", "cosmo", "The name of dataset [cosmo]")
-flags.DEFINE_string("datafile", "data/cosmo_primary_64_1k_train.npy", "Input data file for cosmo")
-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_float("flip_labels", 0, "Probability of flipping labels [0]")
-flags.DEFINE_integer("z_dim", 100, "Dimension of noise vector z [100]")
-flags.DEFINE_integer("nd_layers", 4, "Number of discriminator convolutional layers. [4]")
-flags.DEFINE_integer("ng_layers", 4, "Number of generator conv_T layers. [4]")
-flags.DEFINE_integer("gf_dim", 64, "Dimension of gen filters in last conv layer. [64]")
-flags.DEFINE_integer("df_dim", 64, "Dimension of discrim filters in first conv layer. [64]")
-flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
-flags.DEFINE_integer("output_size", 64, "The size of the output images to produce [64]")
-flags.DEFINE_integer("c_dim", 1, "Dimension of image color. [1]")
-flags.DEFINE_string("data_format", "NHWC", "data format [NHWC]")
-flags.DEFINE_boolean("transpose_matmul_b", False, "Transpose matmul B matrix for performance [False]")
-flags.DEFINE_string("checkpoint_dir", "checkpoints", "Directory name to save the checkpoints [checkpoint]")
-flags.DEFINE_string("experiment", "run_0", "Tensorboard run directory name [run_0]")
-flags.DEFINE_boolean("save_every_step", False, "Save a checkpoint after every step [False]")
-flags.DEFINE_boolean("verbose", True, "print loss on every step [False]")
-flags.DEFINE_integer("num_gpus", -1, "Number of GPUs to use to train GAN. [-1]. -1 means use all available GPUs.")
-config = flags.FLAGS
-
-
-def main(**config_kwargs):
-    import sys
-    # config_kwargs are the default arguments for everything.
-    # These can be overridden by anything specified on the command-line
-    for k, v in config_kwargs.iteritems():
-        try:
-            config[k].value = v
-        except flags.UnrecognizedFlagError:
-            pass
-
-    # Now override the defaults from anything specified on the command-line
-    args = flags.FLAGS(sys.argv, known_only=False)
-    
-    pprint.PrettyPrinter().pprint({k: config[k].value for k in config.__flags })
-    train.train_dcgan(get_data(), config)
-
-
-def get_data():
-    data = np.load(config.datafile, mmap_mode='r')
-
-    if config.data_format == 'NHWC':
-        data = np.expand_dims(data, axis=-1)
-    else:  # 'NCHW'
-        data = np.expand_dims(data, axis=1)
-
-    return data
-
-
-if __name__ == '__main__':
-    tf.app.run()
+import tensorflow as tf
+import train
+import numpy as np
+import pprint
+
+flags = tf.app.flags
+flags.DEFINE_string("dataset", "cosmo", "The name of dataset [cosmo]")
+flags.DEFINE_string("datafile", "data/cosmo_primary_64_1k_train.npy", "Input data file for cosmo")
+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_float("flip_labels", 0, "Probability of flipping labels [0]")
+flags.DEFINE_integer("z_dim", 100, "Dimension of noise vector z [100]")
+flags.DEFINE_integer("nd_layers", 4, "Number of discriminator convolutional layers. [4]")
+flags.DEFINE_integer("ng_layers", 4, "Number of generator conv_T layers. [4]")
+flags.DEFINE_integer("gf_dim", 64, "Dimension of gen filters in last conv layer. [64]")
+flags.DEFINE_integer("df_dim", 64, "Dimension of discrim filters in first conv layer. [64]")
+flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
+flags.DEFINE_integer("output_size", 64, "The size of the output images to produce [64]")
+flags.DEFINE_integer("c_dim", 1, "Dimension of image color. [1]")
+flags.DEFINE_string("data_format", "NHWC", "data format [NHWC]")
+flags.DEFINE_boolean("transpose_matmul_b", False, "Transpose matmul B matrix for performance [False]")
+flags.DEFINE_string("checkpoint_dir", "checkpoints", "Directory name to save the checkpoints [checkpoint]")
+flags.DEFINE_string("experiment", "run_0", "Tensorboard run directory name [run_0]")
+flags.DEFINE_boolean("save_every_step", False, "Save a checkpoint after every step [False]")
+flags.DEFINE_boolean("verbose", True, "print loss on every step [False]")
+flags.DEFINE_integer("num_gpus", -1, "Number of GPUs to use to train GAN. [-1]. -1 means use all available GPUs.")
+config = flags.FLAGS
+
+
+def main(**config_kwargs):
+    import sys
+    # config_kwargs are the default arguments for everything.
+    # These can be overridden by anything specified on the command-line
+    for k, v in config_kwargs.iteritems():
+        try:
+            config[k].value = v
+        except flags.UnrecognizedFlagError:
+            pass
+
+    # Now override the defaults from anything specified on the command-line
+    args = flags.FLAGS(sys.argv, known_only=False)
+    
+    pprint.PrettyPrinter().pprint({k: config[k].value for k in config.__flags })
+    train.train_dcgan(get_data(), config)
+
+
+def get_data():
+    data = np.load(config.datafile, mmap_mode='r')
+
+    if config.data_format == 'NHWC':
+        data = np.expand_dims(data, axis=-1)
+    else:  # 'NCHW'
+        data = np.expand_dims(data, axis=1)
+
+    return data
+
+
+if __name__ == '__main__':
+    tf.app.run()
diff --git a/networks/models/ops.py b/networks/models/ops.py
index 30e9919..1c72a2d 100644
--- a/networks/models/ops.py
+++ b/networks/models/ops.py
@@ -1,142 +1,142 @@
-from collections import defaultdict
-import itertools
-from pprint import pprint
-import sys
-
-import tensorflow as tf
-
-
-def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, transpose_b=False):
-    """Make and return a matrix multiply (input_ * w) + b(ias) operation.
-    
-    """
-
-    shape = input_.get_shape().as_list()
-    if not transpose_b:
-        w_shape = [shape[1], output_size]
-    else:
-        w_shape = [output_size, shape[1]]
-
-    with tf.variable_scope(scope or "linear"):
-        matrix = tf.get_variable('w', w_shape, tf.float32,
-                                 tf.random_normal_initializer(stddev=stddev))
-        bias = tf.get_variable('b', [output_size],
-                               initializer=tf.constant_initializer(bias_start))
-
-        return tf.matmul(input_, matrix, transpose_b=transpose_b) + bias
-
-
-def conv2d(input_, out_channels, data_format, kernel=5, stride=2, stddev=0.02, name="conv2d"):
-    """Make and return a 2D convolution + bias operation.
-    """
-    if data_format == "NHWC":
-        in_channels = input_.get_shape()[-1]
-        strides = [1, stride, stride, 1]
-    else:  # NCHW
-        # WARNING: These strides are probably broken
-        #   https://www.tensorflow.org/api_docs/python/tf/nn/conv2d:
-        #    "Must have strides[0] = strides[3] = 1"
-        in_channels = input_.get_shape()[1]
-        strides = [1, 1, stride, stride]
-
-    with tf.variable_scope(name):
-        # By default, our convolutional mask is a 5x5 filter.
-        w = tf.get_variable('w', [kernel, kernel, in_channels, out_channels],
-                            initializer=tf.truncated_normal_initializer(stddev=stddev))
-        conv = tf.nn.conv2d(input_, w, strides=strides, padding='SAME', data_format=data_format)
-
-        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
-        conv = tf.reshape(tf.nn.bias_add(conv, biases, data_format=data_format), conv.get_shape())
-
-        return conv
-
-
-def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stddev=0.02,
-                     name="conv2d_transpose"):
-
-    if data_format == "NHWC":
-        in_channels = input_.get_shape()[-1]
-        out_channels = output_shape[-1]
-        strides = [1, stride, stride, 1]
-    else:
-        # WARNING: These strides are probably broken
-        #   https://www.tensorflow.org/api_docs/python/tf/nn/conv2d:
-        #    "Must have strides[0] = strides[3] = 1"
-        in_channels = input_.get_shape()[1]
-        out_channels = output_shape[1]
-        strides = [1, 1, stride, stride]
-
-    with tf.variable_scope(name):
-        w = tf.get_variable('w', [kernel, kernel, out_channels, in_channels],
-                            initializer=tf.random_normal_initializer(stddev=stddev))
-
-        deconv = tf.nn.conv2d_transpose(input_, w, output_shape=output_shape, strides=strides, data_format=data_format)
-
-        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
-        deconv = tf.reshape(tf.nn.bias_add(deconv, biases, data_format=data_format), deconv.get_shape())
-
-        return deconv
-
-
-def lrelu(x, alpha=0.2, name="lrelu"):
-    with tf.name_scope(name):
-        return tf.maximum(x, alpha*x)
-
-
-def average_gradients(tower_grads):
-    """Taken from https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py
-
-    Calculate the average gradient for each shared variable across all towers.
-
-    Note that this function provides a synchronization point across all towers.
-    
-    Args:
-    tower_grads: List of lists of (gradient, variable) tuples. The outer list
-      is over individual gradients. The inner list is over the gradient
-      calculation for each tower.
-    Returns:
-     List of pairs of (gradient, variable) where the gradient has been averaged
-     across all towers.
-    """
-    # The variables we get might not have a match from each tower
-    #   or the variable order from each tower may not be the same.
-    # So, we need to do our own grouping here.
-    grouped_vars = defaultdict(list)
-    for grad_and_var in itertools.chain(*tower_grads):
-        # Remove just the "tower_N/" part of the variable name.
-        grouped_var_name = '/'.join(grad_and_var[1].name.split("/")[1:])
-        grouped_vars[grouped_var_name].append(grad_and_var)
-    
-    average_grads = []
-    for grad_and_vars in grouped_vars.itervalues():
-        # Note that each grad_and_vars looks like the following:
-        #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
-        grads = []
-        for g, _ in grad_and_vars:
-            if g is None:
-                continue
-            # Add 0 dimension to the gradients to represent the tower.
-            expanded_g = tf.expand_dims(g, 0)
-
-            # Append on a 'tower' dimension which we will average over below.
-            grads.append(expanded_g)
-
-        # Average over the 'tower' dimension.
-        grad = tf.concat(axis=0, values=grads)
-        grad = tf.reduce_mean(grad, 0)
-
-        # Keep in mind that the Variables are redundant because they are shared
-        # across towers. So .. we will just return the first tower's pointer to
-        # the Variable.
-        # FIXME: This may not be true for the cosmoGAN towers
-        v = grad_and_vars[0][1]
-        grad_and_var = (grad, v)
-        average_grads.append(grad_and_var)
-    
-    return average_grads
-    
-
-def get_available_gpus():
-    from tensorflow.python.client import device_lib
-    local_device_protos = device_lib.list_local_devices()
-    return [x.name for x in local_device_protos if x.device_type == 'GPU']
+from collections import defaultdict
+import itertools
+from pprint import pprint
+import sys
+
+import tensorflow as tf
+
+
+def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, transpose_b=False):
+    """Make and return a matrix multiply (input_ * w) + b(ias) operation.
+    
+    """
+
+    shape = input_.get_shape().as_list()
+    if not transpose_b:
+        w_shape = [shape[1], output_size]
+    else:
+        w_shape = [output_size, shape[1]]
+
+    with tf.variable_scope(scope or "linear"):
+        matrix = tf.get_variable('w', w_shape, tf.float32,
+                                 tf.random_normal_initializer(stddev=stddev))
+        bias = tf.get_variable('b', [output_size],
+                               initializer=tf.constant_initializer(bias_start))
+
+        return tf.matmul(input_, matrix, transpose_b=transpose_b) + bias
+
+
+def conv2d(input_, out_channels, data_format, kernel=5, stride=2, stddev=0.02, name="conv2d"):
+    """Make and return a 2D convolution + bias operation.
+    """
+    if data_format == "NHWC":
+        in_channels = input_.get_shape()[-1]
+        strides = [1, stride, stride, 1]
+    else:  # NCHW
+        # WARNING: These strides are probably broken
+        #   https://www.tensorflow.org/api_docs/python/tf/nn/conv2d:
+        #    "Must have strides[0] = strides[3] = 1"
+        in_channels = input_.get_shape()[1]
+        strides = [1, 1, stride, stride]
+
+    with tf.variable_scope(name):
+        # By default, our convolutional mask is a 5x5 filter.
+        w = tf.get_variable('w', [kernel, kernel, in_channels, out_channels],
+                            initializer=tf.truncated_normal_initializer(stddev=stddev))
+        conv = tf.nn.conv2d(input_, w, strides=strides, padding='SAME', data_format=data_format)
+
+        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
+        conv = tf.reshape(tf.nn.bias_add(conv, biases, data_format=data_format), conv.get_shape())
+
+        return conv
+
+
+def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stddev=0.02,
+                     name="conv2d_transpose"):
+
+    if data_format == "NHWC":
+        in_channels = input_.get_shape()[-1]
+        out_channels = output_shape[-1]
+        strides = [1, stride, stride, 1]
+    else:
+        # WARNING: These strides are probably broken
+        #   https://www.tensorflow.org/api_docs/python/tf/nn/conv2d:
+        #    "Must have strides[0] = strides[3] = 1"
+        in_channels = input_.get_shape()[1]
+        out_channels = output_shape[1]
+        strides = [1, 1, stride, stride]
+
+    with tf.variable_scope(name):
+        w = tf.get_variable('w', [kernel, kernel, out_channels, in_channels],
+                            initializer=tf.random_normal_initializer(stddev=stddev))
+
+        deconv = tf.nn.conv2d_transpose(input_, w, output_shape=output_shape, strides=strides, data_format=data_format)
+
+        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
+        deconv = tf.reshape(tf.nn.bias_add(deconv, biases, data_format=data_format), deconv.get_shape())
+
+        return deconv
+
+
+def lrelu(x, alpha=0.2, name="lrelu"):
+    with tf.name_scope(name):
+        return tf.maximum(x, alpha*x)
+
+
+def average_gradients(tower_grads):
+    """Taken from https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py
+
+    Calculate the average gradient for each shared variable across all towers.
+
+    Note that this function provides a synchronization point across all towers.
+    
+    Args:
+    tower_grads: List of lists of (gradient, variable) tuples. The outer list
+      is over individual gradients. The inner list is over the gradient
+      calculation for each tower.
+    Returns:
+     List of pairs of (gradient, variable) where the gradient has been averaged
+     across all towers.
+    """
+    # The variables we get might not have a match from each tower
+    #   or the variable order from each tower may not be the same.
+    # So, we need to do our own grouping here.
+    grouped_vars = defaultdict(list)
+    for grad_and_var in itertools.chain(*tower_grads):
+        # Remove just the "tower_N/" part of the variable name.
+        grouped_var_name = '/'.join(grad_and_var[1].name.split("/")[1:])
+        grouped_vars[grouped_var_name].append(grad_and_var)
+    
+    average_grads = []
+    for grad_and_vars in grouped_vars.itervalues():
+        # Note that each grad_and_vars looks like the following:
+        #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
+        grads = []
+        for g, _ in grad_and_vars:
+            if g is None:
+                continue
+            # Add 0 dimension to the gradients to represent the tower.
+            expanded_g = tf.expand_dims(g, 0)
+
+            # Append on a 'tower' dimension which we will average over below.
+            grads.append(expanded_g)
+
+        # Average over the 'tower' dimension.
+        grad = tf.concat(axis=0, values=grads)
+        grad = tf.reduce_mean(grad, 0)
+
+        # Keep in mind that the Variables are redundant because they are shared
+        # across towers. So .. we will just return the first tower's pointer to
+        # the Variable.
+        # FIXME: This may not be true for the cosmoGAN towers
+        v = grad_and_vars[0][1]
+        grad_and_var = (grad, v)
+        average_grads.append(grad_and_var)
+    
+    return average_grads
+    
+
+def get_available_gpus():
+    from tensorflow.python.client import device_lib
+    local_device_protos = device_lib.list_local_devices()
+    return [x.name for x in local_device_protos if x.device_type == 'GPU']
diff --git a/networks/models/train.py b/networks/models/train.py
index 96e85b8..79082fe 100644
--- a/networks/models/train.py
+++ b/networks/models/train.py
@@ -1,79 +1,79 @@
-import os
-import sys
-import time
-import numpy as np
-import tensorflow as tf
-import dcgan
-from utils import save_checkpoint, load_checkpoint
-
-
-def train_dcgan(data, config):
-    training_graph = tf.Graph()
-
-    with training_graph.as_default(), tf.device("/cpu:0"):
-        gan = dcgan.DCGAN(output_size=config.output_size,
-                          batch_size=config.batch_size,
-                          nd_layers=config.nd_layers,
-                          ng_layers=config.ng_layers,
-                          df_dim=config.df_dim,
-                          gf_dim=config.gf_dim,
-                          c_dim=config.c_dim,
-                          z_dim=config.z_dim,
-                          flip_labels=config.flip_labels,
-                          data_format=config.data_format,
-                          transpose_b=config.transpose_matmul_b,
-                          num_gpus=config.num_gpus,
-                          learning_rate=config.learning_rate,
-                          beta1=config.beta1)
-
-        save_every_step = True if config.save_every_step == 'True' else False
-
-        g_grads_vars, d_grads_vars = gan.training_graph()
-        update_op = gan.apply_gradients(g_grads_vars, d_grads_vars)
-
-        checkpoint_dir = os.path.join(config.checkpoint_dir, config.experiment)
-
-        with tf.Session() as sess:
-            writer = tf.summary.FileWriter('./logs/'+config.experiment+'/train', sess.graph)
-
-            init_op = tf.global_variables_initializer()
-            sess.run(init_op)
-
-            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, step=save_every_step)
-
-            epoch = sess.run(gan.increment_epoch)
-            start_time = time.time()
-            for epoch in range(epoch, epoch + config.epoch):
-
-                perm = np.random.permutation(data.shape[0])
-                num_batches = data.shape[0] // config.batch_size
-
-                for idx in range(0, num_batches):
-                    batch_images = data[perm[idx*config.batch_size:(idx+1)*config.batch_size]]
-
-                    _, g_sum, d_sum = sess.run([update_op, gan.g_summary, gan.d_summary],
-                                               feed_dict={gan.images: batch_images})
-
-                    global_step = gan.global_step.eval()
-                    writer.add_summary(g_sum, global_step)
-                    writer.add_summary(d_sum, global_step)
-
-                    if save_every_step:
-                        save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, global_step, step=True)
-
-                    if config.verbose:
-                        errD_fake = gan.d_loss_fake.eval()
-                        errD_real = gan.d_loss_real.eval({gan.images: batch_images})
-                        errG = gan.g_loss.eval()
-
-                        print ("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f"
-                               % (epoch, idx, num_batches, time.time() - start_time, errD_fake + errD_real, errG))
-
-                    elif global_step % 100 == 0:
-                        print "Epoch: [%2d] Step: [%4d/%4d] time: %4.4f" % (epoch, idx, num_batches, time.time() - start_time)
-
-                # save a checkpoint every epoch
-                save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, epoch)
-                sess.run(gan.increment_epoch)
-
-
+import os
+import sys
+import time
+import numpy as np
+import tensorflow as tf
+import dcgan
+from utils import save_checkpoint, load_checkpoint
+
+
+def train_dcgan(data, config):
+    training_graph = tf.Graph()
+
+    with training_graph.as_default(), tf.device("/cpu:0"):
+        gan = dcgan.DCGAN(output_size=config.output_size,
+                          batch_size=config.batch_size,
+                          nd_layers=config.nd_layers,
+                          ng_layers=config.ng_layers,
+                          df_dim=config.df_dim,
+                          gf_dim=config.gf_dim,
+                          c_dim=config.c_dim,
+                          z_dim=config.z_dim,
+                          flip_labels=config.flip_labels,
+                          data_format=config.data_format,
+                          transpose_b=config.transpose_matmul_b,
+                          num_gpus=config.num_gpus,
+                          learning_rate=config.learning_rate,
+                          beta1=config.beta1)
+
+        save_every_step = True if config.save_every_step == 'True' else False
+
+        g_grads_vars, d_grads_vars = gan.training_graph()
+        update_op = gan.apply_gradients(g_grads_vars, d_grads_vars)
+
+        checkpoint_dir = os.path.join(config.checkpoint_dir, config.experiment)
+
+        with tf.Session() as sess:
+            writer = tf.summary.FileWriter('./logs/'+config.experiment+'/train', sess.graph)
+
+            init_op = tf.global_variables_initializer()
+            sess.run(init_op)
+
+            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, step=save_every_step)
+
+            epoch = sess.run(gan.increment_epoch)
+            start_time = time.time()
+            for epoch in range(epoch, epoch + config.epoch):
+
+                perm = np.random.permutation(data.shape[0])
+                num_batches = data.shape[0] // config.batch_size
+
+                for idx in range(0, num_batches):
+                    batch_images = data[perm[idx*config.batch_size:(idx+1)*config.batch_size]]
+
+                    _, g_sum, d_sum = sess.run([update_op, gan.g_summary, gan.d_summary],
+                                               feed_dict={gan.images: batch_images})
+
+                    global_step = gan.global_step.eval()
+                    writer.add_summary(g_sum, global_step)
+                    writer.add_summary(d_sum, global_step)
+
+                    if save_every_step:
+                        save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, global_step, step=True)
+
+                    if config.verbose:
+                        errD_fake = gan.d_loss_fake.eval()
+                        errD_real = gan.d_loss_real.eval({gan.images: batch_images})
+                        errG = gan.g_loss.eval()
+
+                        print ("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f"
+                               % (epoch, idx, num_batches, time.time() - start_time, errD_fake + errD_real, errG))
+
+                    elif global_step % 100 == 0:
+                        print "Epoch: [%2d] Step: [%4d/%4d] time: %4.4f" % (epoch, idx, num_batches, time.time() - start_time)
+
+                # save a checkpoint every epoch
+                save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, epoch)
+                sess.run(gan.increment_epoch)
+
+
diff --git a/networks/models/utils.py b/networks/models/utils.py
index 86b80f9..ed8e81b 100644
--- a/networks/models/utils.py
+++ b/networks/models/utils.py
@@ -1,41 +1,41 @@
-import os
-import tensorflow as tf
-
-def save_checkpoint(sess, saver, tag, checkpoint_dir, counter, step=False):
-
-    model_name = tag + '.model-'
-    if step:
-        model_name += 'step'
-    else:
-        model_name += 'epoch'
-
-    if not os.path.exists(checkpoint_dir):
-        os.makedirs(checkpoint_dir)
-
-    saver.save(sess, os.path.join(checkpoint_dir, model_name), global_step=counter)
-
-def load_checkpoint(sess, saver, tag, checkpoint_dir, counter=None, step=False):
-    print(" [*] Reading checkpoints...")
-
-    if step:
-        counter_name = 'step'
-    else:
-        counter_name = 'epoch'
-
-    ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
-    if ckpt and ckpt.model_checkpoint_path:
-        ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
-
-        if not counter==None:
-            ckpt_name_epoch = ckpt_name[:ckpt_name.find(counter_name)] + counter_name + '-%i'%counter
-            if os.path.exists(os.path.join(checkpoint_dir, ckpt_name_epoch+'.index')):
-                ckpt_name = ckpt_name_epoch
-            else:
-                print("Checkpoint for ", counter_name , counter_name, "doesn't exist. Using latest checkpoint instead!")
-
-        saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
-        print(" [*] Success to read {}".format(ckpt_name))
-        return True
-    else:
-        print(" [*] Failed to find a checkpoint")
-        return False
+import os
+import tensorflow as tf
+
+def save_checkpoint(sess, saver, tag, checkpoint_dir, counter, step=False):
+
+    model_name = tag + '.model-'
+    if step:
+        model_name += 'step'
+    else:
+        model_name += 'epoch'
+
+    if not os.path.exists(checkpoint_dir):
+        os.makedirs(checkpoint_dir)
+
+    saver.save(sess, os.path.join(checkpoint_dir, model_name), global_step=counter)
+
+def load_checkpoint(sess, saver, tag, checkpoint_dir, counter=None, step=False):
+    print(" [*] Reading checkpoints...")
+
+    if step:
+        counter_name = 'step'
+    else:
+        counter_name = 'epoch'
+
+    ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
+    if ckpt and ckpt.model_checkpoint_path:
+        ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
+
+        if not counter==None:
+            ckpt_name_epoch = ckpt_name[:ckpt_name.find(counter_name)] + counter_name + '-%i'%counter
+            if os.path.exists(os.path.join(checkpoint_dir, ckpt_name_epoch+'.index')):
+                ckpt_name = ckpt_name_epoch
+            else:
+                print("Checkpoint for ", counter_name , counter_name, "doesn't exist. Using latest checkpoint instead!")
+
+        saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
+        print(" [*] Success to read {}".format(ckpt_name))
+        return True
+    else:
+        print(" [*] Failed to find a checkpoint")
+        return False
diff --git a/networks/run_dcgan.py b/networks/run_dcgan.py
index 288b858..14ed318 100644
--- a/networks/run_dcgan.py
+++ b/networks/run_dcgan.py
@@ -1,55 +1,55 @@
-import os
-
-# These will be the default parameters for the DCGAN
-config_kwargs = dict(
-    datafile = 'data/cosmogan_maps_256_8k_1.npy',
-    output_size = 256,
-    epoch = 50,
-    flip_labels = 0.01,
-    batch_size = 128,
-    z_dim = 64,
-    nd_layers = 4,
-    ng_layers = 4,
-    gf_dim = 64,
-    df_dim = 64,
-    save_every_step = 'False',
-    data_format = 'NCHW',
-    transpose_matmul_b = False,
-    verbose = 'True',
-)
-
-config_kwargs['experiment'] = ('cosmo_myExp_batchSize%i_flipLabel%0.3f_nd%i_ng%i_gfdim%i_dfdim%i_zdim%i' %
-                               (config_kwargs['batch_size'],
-                                config_kwargs['flip_labels'],
-                                config_kwargs['nd_layers'],
-                                config_kwargs['ng_layers'],
-                                config_kwargs['gf_dim'],
-                                config_kwargs['df_dim'],
-                                config_kwargs['z_dim']))
-
-
-if __name__ == "__main__":
-    import models.main
-    import sys
-
-    # The default is to log stdout to an output file.
-    # Only if the user gives the "--stdout" command line option
-    #   should we display stdout to the real stdout.
-    _old_stdout = sys.stdout
-    _new_stdout = sys.stdout
-
-    if "--stdout" in sys.argv:
-        sys.argv.remove("--stdout")
-
-    else:        
-        if not os.path.isdir('output'):
-            os.mkdir('output')
-            
-        _new_stdout = open(os.path.join("output", config_kwargs['experiment']), 'wb')
-        
-
-    try:
-        sys.stdout = _new_stdout
-        models.main.main(**config_kwargs)
-    finally:
-        sys.stdout = _old_stdout
+import os
+
+# These will be the default parameters for the DCGAN
+config_kwargs = dict(
+    datafile = 'data/cosmogan_maps_256_8k_1.npy',
+    output_size = 256,
+    epoch = 50,
+    flip_labels = 0.01,
+    batch_size = 128,
+    z_dim = 64,
+    nd_layers = 4,
+    ng_layers = 4,
+    gf_dim = 64,
+    df_dim = 64,
+    save_every_step = 'False',
+    data_format = 'NCHW',
+    transpose_matmul_b = False,
+    verbose = 'True',
+)
+
+config_kwargs['experiment'] = ('cosmo_myExp_batchSize%i_flipLabel%0.3f_nd%i_ng%i_gfdim%i_dfdim%i_zdim%i' %
+                               (config_kwargs['batch_size'],
+                                config_kwargs['flip_labels'],
+                                config_kwargs['nd_layers'],
+                                config_kwargs['ng_layers'],
+                                config_kwargs['gf_dim'],
+                                config_kwargs['df_dim'],
+                                config_kwargs['z_dim']))
+
+
+if __name__ == "__main__":
+    import models.main
+    import sys
+
+    # The default is to log stdout to an output file.
+    # Only if the user gives the "--stdout" command line option
+    #   should we display stdout to the real stdout.
+    _old_stdout = sys.stdout
+    _new_stdout = sys.stdout
+
+    if "--stdout" in sys.argv:
+        sys.argv.remove("--stdout")
+
+    else:        
+        if not os.path.isdir('output'):
+            os.mkdir('output')
+            
+        _new_stdout = open(os.path.join("output", config_kwargs['experiment']), 'wb')
+        
+
+    try:
+        sys.stdout = _new_stdout
+        models.main.main(**config_kwargs)
+    finally:
+        sys.stdout = _old_stdout