import os, sys

sys.path.append(os.getcwd())

import time

import functools

import numpy as np

import tensorflow as tf

import sklearn.datasets

import tflib as lib

import tflib.ops.linear

import tflib.ops.conv2d

import tflib.ops.batchnorm

import tflib.ops.deconv2d

import tflib.save_images

import tflib.small_imagenet

import tflib.ops.layernorm

import tflib.plot

# Download 64x64 ImageNet at http://image-net.org/small/download.php and

# fill in the path to the extracted files here!

DATA_DIR = ''

if len(DATA_DIR) == 0:

raise Exception('Please specify path to data directory in gan_64x64.py!')

MODE = 'wgan-gp' # dcgan, wgan, wgan-gp, lsgan

DIM = 64 # Model dimensionality

CRITIC_ITERS = 5 # How many iterations to train the critic for

N_GPUS = 1 # Number of GPUs

BATCH_SIZE = 64 # Batch size. Must be a multiple of N_GPUS

ITERS = 200000 # How many iterations to train for

LAMBDA = 10 # Gradient penalty lambda hyperparameter

OUTPUT_DIM = 64*64*3 # Number of pixels in each iamge

lib.print_model_settings(locals().copy())

def GeneratorAndDiscriminator():

raise Exception('You must choose an architecture!')

DEVICES = ['/gpu:{}'.format(i) for i in xrange(N_GPUS)]

def LeakyReLU(x, alpha=0.2):

return tf.maximum(alpha*x, x)

def ReLULayer(name, n_in, n_out, inputs):

output = lib.ops.linear.Linear(name+'.Linear', n_in, n_out, inputs, initialization='he')

return tf.nn.relu(output)

def LeakyReLULayer(name, n_in, n_out, inputs):

output = lib.ops.linear.Linear(name+'.Linear', n_in, n_out, inputs, initialization='he')

return LeakyReLU(output)

def Normalize(name, axes, inputs):

return lib.ops.batchnorm.Batchnorm(name,axes,inputs,fused=True)

def pixcnn_gated_nonlinearity(a, b):

return tf.sigmoid(a) * tf.tanh(b)

def SubpixelConv2D(*args, **kwargs):

return output

def ConvMeanPool(name, input_dim, output_dim, filter_size, inputs, he_init=True, biases=True):

output = lib.ops.conv2d.Conv2D(name, input_dim, output_dim, filter_size, inputs, he_init=he_init, biases=biases)

output = tf.add_n([output[:,:,::2,::2], output[:,:,1::2,::2], output[:,:,::2,1::2], output[:,:,1::2,1::2]]) / 4.

return output

def MeanPoolConv(name, input_dim, output_dim, filter_size, inputs, he_init=True, biases=True):

return output

def UpsampleConv(name, input_dim, output_dim, filter_size, inputs, he_init=True, biases=True):

return output

def BottleneckResidualBlock(name, input_dim, output_dim, filter_size, inputs, resample=None, he_init=True):

return shortcut + (0.3*output)

def ResidualBlock(name, input_dim, output_dim, filter_size, inputs, resample=None, he_init=True):

return shortcut + output

# ! Generators

def GoodGenerator(n_samples, noise=None, dim=DIM, nonlinearity=tf.nn.relu):

return tf.reshape(output, [-1, OUTPUT_DIM])

def FCGenerator(n_samples, noise=None, FC_DIM=512):

return output

def DCGANGenerator(n_samples, noise=None, dim=DIM, bn=True, nonlinearity=tf.nn.relu):

return tf.reshape(output, [-1, OUTPUT_DIM])

def WGANPaper_CrippledDCGANGenerator(n_samples, noise=None, dim=DIM):

return tf.reshape(output, [-1, OUTPUT_DIM])

def ResnetGenerator(n_samples, noise=None, dim=DIM):

return tf.reshape(output, [-1, OUTPUT_DIM])

def MultiplicativeDCGANGenerator(n_samples, noise=None, dim=DIM, bn=True):

return tf.reshape(output, [-1, OUTPUT_DIM])

# ! Discriminators

def GoodDiscriminator(inputs, dim=DIM):

return tf.reshape(output, [-1])

def MultiplicativeDCGANDiscriminator(inputs, dim=DIM, bn=True):

return tf.reshape(output, [-1])

def ResnetDiscriminator(inputs, dim=DIM):

return tf.reshape(output / 5., [-1])

def FCDiscriminator(inputs, FC_DIM=512, n_layers=3):

return tf.reshape(output, [-1])

def DCGANDiscriminator(inputs, dim=DIM, bn=True, nonlinearity=LeakyReLU):

return tf.reshape(output, [-1])

Generator, Discriminator = GeneratorAndDiscriminator()

with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as session:

all_real_data_conv = tf.placeholder(tf.int32, shape=[BATCH_SIZE, 3, 64, 64])

if tf.__version__.startswith('1.'):

split_real_data_conv = tf.split(all_real_data_conv, len(DEVICES))

else:

split_real_data_conv = tf.split(0, len(DEVICES), all_real_data_conv)

gen_costs, disc_costs = [],[]

for device_index, (device, real_data_conv) in enumerate(zip(DEVICES, split_real_data_conv)):

with tf.device(device):

real_data = tf.reshape(2*((tf.cast(real_data_conv, tf.float32)/255.)-.5), [BATCH_SIZE/len(DEVICES), OUTPUT_DIM])

fake_data = Generator(BATCH_SIZE/len(DEVICES))

disc_real = Discriminator(real_data)

disc_fake = Discriminator(fake_data)

if MODE == 'wgan':

gen_cost = -tf.reduce_mean(disc_fake)

disc_cost = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)

elif MODE == 'wgan-gp':

gen_cost = -tf.reduce_mean(disc_fake)

disc_cost = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)

alpha = tf.random_uniform(

shape=[BATCH_SIZE/len(DEVICES),1],

minval=0.,

maxval=1.

)

differences = fake_data - real_data

interpolates = real_data + (alpha*differences)

gradients = tf.gradients(Discriminator(interpolates), [interpolates])[0]

slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))

gradient_penalty = tf.reduce_mean((slopes-1.)**2)

disc_cost += LAMBDA*gradient_penalty

elif MODE == 'dcgan':

try: # tf pre-1.0 (bottom) vs 1.0 (top)

gen_cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=disc_fake,

labels=tf.ones_like(disc_fake)))

disc_cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=disc_fake,

labels=tf.zeros_like(disc_fake)))

disc_cost += tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=disc_real,

labels=tf.ones_like(disc_real)))

except Exception as e:

gen_cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(disc_fake, tf.ones_like(disc_fake)))

disc_cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(disc_fake, tf.zeros_like(disc_fake)))

disc_cost += tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(disc_real, tf.ones_like(disc_real)))

disc_cost /= 2.

elif MODE == 'lsgan':

gen_cost = tf.reduce_mean((disc_fake - 1)**2)

disc_cost = (tf.reduce_mean((disc_real - 1)**2) + tf.reduce_mean((disc_fake - 0)**2))/2.

else:

raise Exception()

gen_costs.append(gen_cost)

disc_costs.append(disc_cost)

gen_cost = tf.add_n(gen_costs) / len(DEVICES)

disc_cost = tf.add_n(disc_costs) / len(DEVICES)

if MODE == 'wgan':

gen_train_op = tf.train.RMSPropOptimizer(learning_rate=5e-5).minimize(gen_cost,

var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)

disc_train_op = tf.train.RMSPropOptimizer(learning_rate=5e-5).minimize(disc_cost,

var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

clip_ops = []

for var in lib.params_with_name('Discriminator'):

clip_bounds = [-.01, .01]

clip_ops.append(tf.assign(var, tf.clip_by_value(var, clip_bounds[0], clip_bounds[1])))

clip_disc_weights = tf.group(*clip_ops)

elif MODE == 'wgan-gp':

gen_train_op = tf.train.AdamOptimizer(learning_rate=1e-4, beta1=0., beta2=0.9).minimize(gen_cost,

var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)

disc_train_op = tf.train.AdamOptimizer(learning_rate=1e-4, beta1=0., beta2=0.9).minimize(disc_cost,

var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

elif MODE == 'dcgan':

gen_train_op = tf.train.AdamOptimizer(learning_rate=2e-4, beta1=0.5).minimize(gen_cost,

var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)

disc_train_op = tf.train.AdamOptimizer(learning_rate=2e-4, beta1=0.5).minimize(disc_cost,

var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

elif MODE == 'lsgan':

gen_train_op = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(gen_cost,

var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)

disc_train_op = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(disc_cost,

var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

else:

raise Exception()

# For generating samples

fixed_noise = tf.constant(np.random.normal(size=(BATCH_SIZE, 128)).astype('float32'))

all_fixed_noise_samples = []

for device_index, device in enumerate(DEVICES):

n_samples = BATCH_SIZE / len(DEVICES)

all_fixed_noise_samples.append(Generator(n_samples, noise=fixed_noise[device_index*n_samples:(device_index+1)*n_samples]))

if tf.__version__.startswith('1.'):

all_fixed_noise_samples = tf.concat(all_fixed_noise_samples, axis=0)

else:

all_fixed_noise_samples = tf.concat(0, all_fixed_noise_samples)

def generate_image(iteration):

samples = session.run(all_fixed_noise_samples)

samples = ((samples+1.)*(255.99/2)).astype('int32')

lib.save_images.save_images(samples.reshape((BATCH_SIZE, 3, 64, 64)), 'samples_{}.png'.format(iteration))

# Dataset iterator

train_gen, dev_gen = lib.small_imagenet.load(BATCH_SIZE, data_dir=DATA_DIR)

def inf_train_gen():

while True:

for (images,) in train_gen():

yield images

# Save a batch of ground-truth samples

_x = inf_train_gen().next()

_x_r = session.run(real_data, feed_dict={real_data_conv: _x[:BATCH_SIZE/N_GPUS]})

_x_r = ((_x_r+1.)*(255.99/2)).astype('int32')

lib.save_images.save_images(_x_r.reshape((BATCH_SIZE/N_GPUS, 3, 64, 64)), 'samples_groundtruth.png')

# Train loop

session.run(tf.initialize_all_variables())

gen = inf_train_gen()

for iteration in xrange(ITERS):

start_time = time.time()

# Train generator

if iteration > 0:

_ = session.run(gen_train_op)

# Train critic

if (MODE == 'dcgan') or (MODE == 'lsgan'):

disc_iters = 1

else:

disc_iters = CRITIC_ITERS

for i in xrange(disc_iters):

_data = gen.next()

_disc_cost, _ = session.run([disc_cost, disc_train_op], feed_dict={all_real_data_conv: _data})

if MODE == 'wgan':

_ = session.run([clip_disc_weights])

lib.plot.plot('train disc cost', _disc_cost)

lib.plot.plot('time', time.time() - start_time)

if iteration % 200 == 199:

t = time.time()

dev_disc_costs = []

for (images,) in dev_gen():

_dev_disc_cost = session.run(disc_cost, feed_dict={all_real_data_conv: images})

dev_disc_costs.append(_dev_disc_cost)

lib.plot.plot('dev disc cost', np.mean(dev_disc_costs))

generate_image(iteration)

if (iteration < 5) or (iteration % 200 == 199):

lib.plot.flush()

lib.plot.tick()