In [14]:
# Make TFGAN models discoverable
import sys
import os
sys.path.append('gan')
sys.path.append('slim')
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from six.moves import xrange
import tensorflow as tf
# TFGAN library
tfgan = tf.contrib.gan
# Shortcuts
queues = tf.contrib.slim.queues
layers = tf.contrib.layers
framework = tf.contrib.framework
tf.reset_default_graph()
In [15]:
def evaluate_tfgan_loss(gan_loss, name=None):
"""Evaluate GAN losses. Used to check that the graph is correct.
Args:
gan_loss: A GANLoss tuple.
name: Optional. If present, append to debug output.
"""
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
with queues.QueueRunners(sess):
gen_loss_np = sess.run(gan_loss.generator_loss)
dis_loss_np = sess.run(gan_loss.discriminator_loss)
if name:
print("%s generator loss: %f" % (name, gen_loss_np))
print("%s discriminator loss: %f" % (name, dis_loss_np))
else:
print("Generator loss: %f" % gen_loss_np)
print("Discriminator loss: %f" % dis_loss_np)
In [16]:
n_step = 64 # Number of time steps for ODE solve
step_size = 0.01 # Step size
n_sample = 4 # Number of samples for input to the generator
stride = n_step // n_sample # stride used to subsample the data
assert stride*n_sample == n_step, "Wrong value for stride"
output_size = n_step + 1 # Size of the generator output
input_size = 2*n_sample # Size of the generator input
# RHS of the ODE
def f(t,a,b):
return a * np.exp(-20*t) + b * np.cos(64*t)
# Euler time integrator
def solve(y0,a,b,h,n):
# Solve ODE with initial condition y0
y = np.empty(n+1, dtype=np.float32)
# Initial condition
y[0] = np.float32(y0)
# Simplified ODE formula
for i in xrange(1,n+1):
y[i] = y[i-1] - 0.05 * y[i-1] + f(i*h,a,b)
return y
# Generate input/output data by solving the ODE
def generate_data(n_data):
inputs = np.empty([n_data, input_size], dtype=np.float32)
outputs = np.empty([inputs.shape[0], output_size], dtype=np.float32)
for i in xrange(inputs.shape[0]):
y0 = np.random.rand() # Initial condition
a = 0.1 * np.random.rand() # Parameters for RHS
b = 0.1 * np.random.rand()
# Time steps
t = np.linspace(0, step_size*(outputs.shape[1]-1), outputs.shape[1])
f_eval = f(t,a,b) # RHS
y = solve(y0,a,b,step_size,n_step) # Solution y(t)
# Subsample solution and RHS
y_stride = y[0:n_step:stride]
f_stride = f_eval[0:n_step:stride]
# Concatenate to the get inputs
inputs[i,:] = np.concatenate((y_stride, f_stride), axis=None)
assert len(inputs[i,:]) == input_size, "inputs does not have the right length"
# Output is the solution
outputs[i,:] = y
assert len(outputs[i,:]) == output_size, "outputs does not have the right length"
assert np.linalg.norm(outputs[i,0:n_step:stride]-inputs[i,:inputs.shape[1]//2]) == 0, "outputs does not have the right length"
rel_noise = 0 # 0.05
inputs = inputs * np.float32(1 + rel_noise * 2 * (np.random.rand(*inputs.shape) - 0.5))
outputs = outputs * np.float32(1 + rel_noise * 2 * (np.random.rand(*outputs.shape) - 0.5))
assert inputs.shape[1] == input_size, "inputs does not have the right shape"
assert outputs.shape[1] == output_size, "outputs does not have the right shape"
assert inputs.shape[0] == n_data, "inputs does not have the right shape"
assert outputs.shape[0] == n_data, "outputs does not have the right shape"
return inputs, outputs
In [17]:
def generator_fn(inputs, is_training=True):
real_input, _ = inputs
assert real_input.shape[1] == input_size, "real_input does not have the right length"
with framework.arg_scope([layers.fully_connected]):
return layers.linear(real_input, output_size)
def discriminator_fn(inputs, conditioning, is_training=True):
_, conditioning_data = conditioning
assert conditioning_data.shape[1] == input_size, "conditioning_data does not have the right length"
with framework.arg_scope([layers.fully_connected]):
net = layers.linear(inputs, 128)
net = tfgan.features.condition_tensor(net, conditioning_data)
net = layers.fully_connected(net, 64, normalizer_fn=layers.layer_norm)
return layers.linear(net, 1)
In [18]:
batch_size = 32
n_data = 128 * batch_size
In [19]:
inputs, real_data = generate_data(n_data)
assert inputs.shape[0] == n_data, "wrong size for inputs"
assert inputs.shape[1] == input_size, "wrong size for inputs"
assert real_data.shape[1] == output_size, "wrong size for real_data"
conditional_gan_model = tfgan.gan_model(
generator_fn=generator_fn,
discriminator_fn=discriminator_fn,
real_data=real_data,
generator_inputs=(inputs,inputs))
In [20]:
plt.figure(figsize=(10, 3), dpi=100)
plt.subplots_adjust(wspace=0.3, hspace=0.5)
plt.subplot(1,3,1)
plt.title("y samples")
plt.plot(inputs[0,:inputs.shape[1]//2],'o')
plt.subplot(1,3,2)
plt.title("RHS samples")
plt.plot(inputs[0,inputs.shape[1]//2:],'s')
plt.subplot(1,3,3)
plt.title("Output")
plt.plot(real_data[0,:])
Out[20]:
In [21]:
# Wasserstein loss (https://arxiv.org/abs/1701.07875) with the
# gradient penalty from the improved Wasserstein loss paper
# (https://arxiv.org/abs/1704.00028).
gan_loss = tfgan.gan_loss(conditional_gan_model,gradient_penalty_weight=1.0)
# Vanilla loss function
# generator_loss_fn=tfgan.losses.minimax_generator_loss,
# discriminator_loss_fn=tfgan.losses.minimax_discriminator_loss)
evaluate_tfgan_loss(gan_loss, "gan loss")
In [22]:
# Create the training optimizers, which calculate gradients and apply updates to weights.
# generator_optimizer = tf.train.AdamOptimizer(0.0009, beta1=0.5)
# discriminator_optimizer = tf.train.AdamOptimizer(0.00009, beta1=0.5)
# generator_optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.999, epsilon=0.01)
# discriminator_optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.999, epsilon=0.01)
# generator_optimizer = tf.train.GradientDescentOptimizer(0.01)
# discriminator_optimizer = tf.train.GradientDescentOptimizer(0.01)
generator_optimizer = tf.train.MomentumOptimizer(learning_rate=0.005,momentum=0.2)
discriminator_optimizer = tf.train.MomentumOptimizer(learning_rate=0.005,momentum=0.2)
# generator_optimizer = tf.train.RMSPropOptimizer(0.001)
# discriminator_optimizer = tf.train.RMSPropOptimizer(0.001)
gan_train_ops = tfgan.gan_train_ops(
conditional_gan_model,
gan_loss,
generator_optimizer,
discriminator_optimizer)
In [23]:
# Data used for training
inputs, outputs = generate_data(n_data)
with tf.variable_scope("Generator", reuse=True):
gan_prediction = conditional_gan_model.generator_fn((inputs, inputs))
In [24]:
global_step = tf.train.get_or_create_global_step()
train_step_fn = tfgan.get_sequential_train_steps()
# Post-processing
loss_values = []
pred = []
# Number of epochs for training
n_train = 10000
with tf.train.SingularMonitoredSession() as sess:
for i in xrange(n_train+1):
cur_loss, _ = train_step_fn(
sess, gan_train_ops, global_step, train_step_kwargs={})
# Append values
loss_values.append((i, cur_loss))
pred.append(sess.run([gan_prediction]))
# Print current loss
if i % (n_train//20) == 0:
print("Loss: %11.3f step: %6d/%6d" % (cur_loss, i, n_train))
In [25]:
plt.title("Training loss")
plt.plot(*zip(*loss_values))
Out[25]:
In [26]:
n_pred = len(pred)
n_plts = 4
n_figs = 16
time_plt = np.array(np.linspace(0,n_pred-1,n_figs), dtype = int)
print(time_plt)
plt.figure(figsize=(11, 3), dpi=90)
plt.subplot(1,2,1)
plt.title("Real data")
plt_sample = inputs.shape[1]//2
plt.plot(np.linspace(0,n_step - n_step//plt_sample,plt_sample,plt_sample),\
inputs[0,:plt_sample],'ks',fillstyle='none')
for k in xrange(n_plts):
plt.plot(outputs[k,:])
plt.subplot(1,2,2)
plt.title("GAN prediction")
for k in xrange(n_plts):
plt.plot(pred[-1][0][k])
plt.figure(figsize=(11, 6), dpi=90)
plt.subplots_adjust(wspace=0.4, hspace=0.4)
for idx in xrange(n_figs):
plt.subplot(4,4,idx+1)
for k in xrange(n_plts):
plt.plot(pred[time_plt[idx]][0][k])
plt.suptitle("Evolution of GAN")
Out[26]:
