diff --git a/.gitignore b/.gitignore index 0d20b64..94487b9 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1 @@ -*.pyc +*.pyc diff --git a/networks/models/dcgan.py b/networks/models/dcgan.py index b90d1a2..df904d5 100644 --- a/networks/models/dcgan.py +++ b/networks/models/dcgan.py @@ -1,83 +1,144 @@ +import sys + import tensorflow as tf -from .ops import linear, conv2d, conv2d_transpose, lrelu +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, +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): + 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 - 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.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.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, + 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]) - with tf.variable_scope("discriminator") as d_scope: - d_prob_real, d_logits_real = self.discriminator(self.images, is_training=True) - - with tf.variable_scope("generator") as g_scope: - 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))) - + # 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", self.d_loss_real), - tf.summary.scalar("loss/fake", self.d_loss_fake), - tf.summary.scalar("loss/d", self.d_loss)]) + 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", self.g_loss)] - if self.data_format == "NHWC": # tf.summary.image is not implemented for NCHW + 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) - 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") 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() 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": @@ -87,75 +148,74 @@ 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() 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) + 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) - g_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1) \ - .minimize(self.g_loss, var_list=self.g_vars) + return d_optim, g_optim - return tf.group(d_optim, g_optim, name="all_optims") + 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)) - 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) + 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) - # 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)) + 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""" - # h0 = lrelu(conv2d(image)) - h0 = lrelu(conv2d(image, self.df_dim, self.data_format, name='h0_conv')) + chain = 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 = 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 @@ -191,3 +251,4 @@ def _check_architecture_consistency(self): 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..41f8d93 100644 --- a/networks/models/main.py +++ b/networks/models/main.py @@ -24,22 +24,37 @@ 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) +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' + 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..1c72a2d 100644 --- a/networks/models/ops.py +++ b/networks/models/ops.py @@ -1,6 +1,15 @@ +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: @@ -12,20 +21,26 @@ def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, transpo 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)) + 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"): +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 + 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) @@ -35,6 +50,7 @@ def conv2d(input_, out_channels, data_format, kernel=5, stride=2, stddev=0.02, n return conv + def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stddev=0.02, name="conv2d_transpose"): @@ -43,6 +59,9 @@ def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stdd 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] @@ -50,15 +69,74 @@ def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stdd 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) + 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 0b52f69..79082fe 100644 --- a/networks/models/train.py +++ b/networks/models/train.py @@ -1,17 +1,17 @@ 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): +def train_dcgan(data, config): training_graph = tf.Graph() - with training_graph.as_default(): - - gan = dcgan.dcgan(output_size=config.output_size, + 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, @@ -21,12 +21,15 @@ def train_dcgan(data, config): z_dim=config.z_dim, flip_labels=config.flip_labels, data_format=config.data_format, - transpose_b=config.transpose_matmul_b) + 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 - gan.training_graph() - update_op = gan.optimizer(config.learning_rate, config.beta1) + 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) @@ -48,7 +51,7 @@ def train_dcgan(data, config): 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], + _, 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() @@ -63,12 +66,14 @@ 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)) + 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/run_dcgan.py b/networks/run_dcgan.py index 05895a1..14ed318 100644 --- a/networks/run_dcgan.py +++ b/networks/run_dcgan.py @@ -1,37 +1,55 @@ 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) + +# 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