def train_vaegan(files,
learning_rate=0.00001,
batch_size=64,
n_epochs=250,
n_examples=10,
input_shape=[218, 178, 3],
crop_shape=[64, 64, 3],
crop_factor=0.8,
n_filters=[100, 100, 100, 100],
n_hidden=None,
n_code=128,
convolutional=True,
variational=True,
filter_sizes=[3, 3, 3, 3],
activation=tf.nn.elu,
ckpt_name="vaegan.ckpt"):
"""Summary
Parameters
----------
files : TYPE
Description
learning_rate : float, optional
Description
batch_size : int, optional
Description
n_epochs : int, optional
Description
n_examples : int, optional
Description
input_shape : list, optional
Description
crop_shape : list, optional
Description
crop_factor : float, optional
Description
n_filters : list, optional
Description
n_hidden : int, optional
Description
n_code : int, optional
Description
convolutional : bool, optional
Description
variational : bool, optional
Description
filter_sizes : list, optional
Description
activation : TYPE, optional
Description
ckpt_name : str, optional
Description
No Longer Returned
------------------
name : TYPE
Description
"""
ae = VAEGAN(input_shape=[None] + crop_shape,
convolutional=convolutional,
variational=variational,
n_filters=n_filters,
n_hidden=n_hidden,
n_code=n_code,
filter_sizes=filter_sizes,
activation=activation)
batch = create_input_pipeline(files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=crop_shape,
crop_factor=crop_factor,
shape=input_shape)
zs = np.random.randn(4, n_code).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
opt_enc = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['loss_enc'],
var_list=[
var_i for var_i in tf.trainable_variables()
if var_i.name.startswith('encoder')
])
opt_gen = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['loss_gen'],
var_list=[
var_i for var_i in tf.trainable_variables()
if var_i.name.startswith('generator')
])
opt_dis = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['loss_dis'],
var_list=[
var_i for var_i in tf.trainable_variables()
if var_i.name.startswith('discriminator')
])
sess = tf.Session()
saver = tf.train.Saver()
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
if os.path.exists(ckpt_name + '.index') or os.path.exists(ckpt_name):
saver.restore(sess, ckpt_name)
print("VAE model restored.")
t_i = 0
batch_i = 0
epoch_i = 0
equilibrium = 0.693
margin = 0.4
n_files = len(files)
test_xs = sess.run(batch) / 255.0
utils.montage(test_xs, 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
if batch_i % (n_files // batch_size) == 0:
batch_i = 0
epoch_i += 1
print('---------- EPOCH:', epoch_i)
batch_i += 1
batch_xs = sess.run(batch) / 255.0
batch_zs = np.random.randn(batch_size, n_code).astype(np.float32)
real_cost, fake_cost, _ = sess.run(
[ae['loss_real'], ae['loss_fake'], opt_enc],
feed_dict={
ae['x']: batch_xs,
ae['gamma']: 0.5
})
real_cost = -np.mean(real_cost)
fake_cost = -np.mean(fake_cost)
print('real:', real_cost, '/ fake:', fake_cost)
gen_update = True
dis_update = True
if real_cost > (equilibrium + margin) or \
fake_cost > (equilibrium + margin):
gen_update = False
if real_cost < (equilibrium - margin) or \
fake_cost < (equilibrium - margin):
dis_update = False
if not (gen_update or dis_update):
gen_update = True
dis_update = True
if gen_update:
sess.run(opt_gen,
feed_dict={
ae['x']: batch_xs,
ae['z_samp']: batch_zs,
ae['gamma']: 0.5
})
if dis_update:
sess.run(opt_dis,
feed_dict={
ae['x']: batch_xs,
ae['z_samp']: batch_zs,
ae['gamma']: 0.5
})
if batch_i % 50 == 0:
# Plot example reconstructions from latent layer
recon = sess.run(ae['x_tilde'], feed_dict={ae['z']: zs})
print('recon:', recon.min(), recon.max())
recon = np.clip(recon / recon.max(), 0, 1)
utils.montage(recon.reshape([-1] + crop_shape),
'imgs/manifold_%08d.png' % t_i)
# Plot example reconstructions
recon = sess.run(ae['x_tilde'], feed_dict={ae['x']: test_xs})
print('recon:', recon.min(), recon.max())
recon = np.clip(recon / recon.max(), 0, 1)
utils.montage(recon.reshape([-1] + crop_shape),
'imgs/reconstruction_%08d.png' % t_i)
t_i += 1
if batch_i % 100 == 0:
# Save the variables to disk.
save_path = saver.save(sess,
ckpt_name,
global_step=batch_i,
write_meta_graph=False)
print("Model saved in file: %s" % save_path)
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()
def train_vae(files,
input_shape,
learning_rate=0.0001,
batch_size=100,
n_epochs=50,
n_examples=10,
crop_shape=[64, 64, 3],
crop_factor=0.8,
n_filters=[100, 100, 100, 100],
n_hidden=256,
n_code=50,
convolutional=True,
variational=True,
filter_sizes=[3, 3, 3, 3],
dropout=True,
keep_prob=0.8,
activation=tf.nn.relu,
img_step=100,
save_step=100,
ckpt_name="vae.ckpt"):
"""General purpose training of a (Variational) (Convolutional) Autoencoder.
Supply a list of file paths to images, and this will do everything else.
Parameters
----------
files : list of strings
List of paths to images.
input_shape : list
Must define what the input image's shape is.
learning_rate : float, optional
Learning rate.
batch_size : int, optional
Batch size.
n_epochs : int, optional
Number of epochs.
n_examples : int, optional
Number of example to use while demonstrating the current training
iteration's reconstruction. Creates a square montage, so make
sure int(sqrt(n_examples))**2 = n_examples, e.g. 16, 25, 36, ... 100.
crop_shape : list, optional
Size to centrally crop the image to.
crop_factor : float, optional
Resize factor to apply before cropping.
n_filters : list, optional
Same as VAE's n_filters.
n_hidden : int, optional
Same as VAE's n_hidden.
n_code : int, optional
Same as VAE's n_code.
convolutional : bool, optional
Use convolution or not.
variational : bool, optional
Use variational layer or not.
filter_sizes : list, optional
Same as VAE's filter_sizes.
dropout : bool, optional
Use dropout or not
keep_prob : float, optional
Percent of keep for dropout.
activation : function, optional
Which activation function to use.
img_step : int, optional
How often to save training images showing the manifold and
reconstruction.
save_step : int, optional
How often to save checkpoints.
ckpt_name : str, optional
Checkpoints will be named as this, e.g. 'model.ckpt'
"""
batch = create_input_pipeline(
files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=crop_shape,
crop_factor=crop_factor,
shape=input_shape)
ae = VAE(
input_shape=[None] + crop_shape,
convolutional=convolutional,
variational=variational,
n_filters=n_filters,
n_hidden=n_hidden,
n_code=n_code,
dropout=dropout,
filter_sizes=filter_sizes,
activation=activation)
# Create a manifold of our inner most layer to show
# example reconstructions. This is one way to see
# what the "embedding" or "latent space" of the encoder
# is capable of encoding, though note that this is just
# a random hyperplane within the latent space, and does not
# encompass all possible embeddings.
zs = np.random.uniform(-1.0, 1.0, [4, n_code]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
optimizer = tf.train.AdamOptimizer(
learning_rate=learning_rate).minimize(ae['cost'])
# We create a session to use the graph
sess = tf.Session()
saver = tf.train.Saver()
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
# This will handle our threaded image pipeline
coord = tf.train.Coordinator()
# Ensure no more changes to graph
tf.get_default_graph().finalize()
# Start up the queues for handling the image pipeline
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
if os.path.exists(ckpt_name + '.index') or os.path.exists(ckpt_name):
saver.restore(sess, ckpt_name)
# Fit all training data
t_i = 0
batch_i = 0
epoch_i = 0
cost = 0
n_files = len(files)
test_xs = sess.run(batch) / 255.0
utils.montage(test_xs, 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
batch_i += 1
batch_xs = sess.run(batch) / 255.0
train_cost = sess.run(
[ae['cost'], optimizer],
feed_dict={
ae['x']: batch_xs,
ae['train']: True,
ae['keep_prob']: keep_prob
})[0]
print(batch_i, train_cost)
cost += train_cost
if batch_i % n_files == 0:
print('epoch:', epoch_i)
print('average cost:', cost / batch_i)
cost = 0
batch_i = 0
epoch_i += 1
if batch_i % img_step == 0:
# Plot example reconstructions from latent layer
recon = sess.run(
ae['y'],
feed_dict={
ae['z']: zs,
ae['train']: False,
ae['keep_prob']: 1.0
})
utils.montage(
recon.reshape([-1] + crop_shape), 'manifold_%08d.png' % t_i)
# Plot example reconstructions
recon = sess.run(
ae['y'],
feed_dict={
ae['x']: test_xs,
ae['train']: False,
ae['keep_prob']: 1.0
})
print('reconstruction (min, max, mean):',
recon.min(), recon.max(), recon.mean())
utils.montage(
recon.reshape([-1] + crop_shape),
'reconstruction_%08d.png' % t_i)
t_i += 1
if batch_i % save_step == 0:
# Save the variables to disk.
saver.save(
sess,
ckpt_name,
global_step=batch_i,
write_meta_graph=False)
except tf.errors.OutOfRangeError:
print('Done.')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()
def test_mnist():
"""Train an autoencoder on MNIST.
This function will train an autoencoder on MNIST and also
save many image files during the training process, demonstrating
the latent space of the inner most dimension of the encoder,
as well as reconstructions of the decoder.
"""
# load MNIST
n_code = 2
mnist = MNIST(split=[0.8, 0.1, 0.1])
ae = VAE(
input_shape=[None, 784],
n_filters=[512, 256],
n_hidden=64,
n_code=n_code,
activation=tf.nn.sigmoid,
convolutional=False,
variational=True)
n_examples = 100
zs = np.random.uniform(-1.0, 1.0, [4, n_code]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
learning_rate = 0.02
optimizer = tf.train.AdamOptimizer(
learning_rate=learning_rate).minimize(ae['cost'])
# We create a session to use the graph
sess = tf.Session()
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
# Fit all training data
t_i = 0
batch_i = 0
batch_size = 200
n_epochs = 10
test_xs = mnist.test.images[:n_examples]
utils.montage(test_xs.reshape((-1, 28, 28)), 'test_xs.png')
for epoch_i in range(n_epochs):
train_i = 0
train_cost = 0
for batch_xs, _ in mnist.train.next_batch(batch_size):
train_cost += sess.run(
[ae['cost'], optimizer],
feed_dict={
ae['x']: batch_xs,
ae['train']: True,
ae['keep_prob']: 1.0
})[0]
train_i += 1
if batch_i % 10 == 0:
# Plot example reconstructions from latent layer
recon = sess.run(
ae['y'],
feed_dict={
ae['z']: zs,
ae['train']: False,
ae['keep_prob']: 1.0
})
utils.montage(
recon.reshape((-1, 28, 28)), 'manifold_%08d.png' % t_i)
# Plot example reconstructions
recon = sess.run(
ae['y'],
feed_dict={
ae['x']: test_xs,
ae['train']: False,
ae['keep_prob']: 1.0
})
utils.montage(
recon.reshape((-1, 28, 28)),
'reconstruction_%08d.png' % t_i)
t_i += 1
batch_i += 1
valid_i = 0
valid_cost = 0
for batch_xs, _ in mnist.valid.next_batch(batch_size):
valid_cost += sess.run(
[ae['cost']],
feed_dict={
ae['x']: batch_xs,
ae['train']: False,
ae['keep_prob']: 1.0
})[0]
valid_i += 1
print('train:', train_cost / train_i, 'valid:', valid_cost / valid_i)
def train_vae(files,
input_shape,
learning_rate=0.0001,
batch_size=100,
n_epochs=50,
n_examples=10,
crop_shape=[64, 64, 3],
crop_factor=0.8,
n_filters=[100, 100, 100, 100],
n_hidden=256,
n_code=50,
convolutional=True,
variational=True,
filter_sizes=[3, 3, 3, 3],
dropout=True,
keep_prob=0.8,
activation=tf.nn.relu,
img_step=100,
save_step=100,
ckpt_name="vae.ckpt"):
"""General purpose training of a (Variational) (Convolutional) Autoencoder.
Supply a list of file paths to images, and this will do everything else.
Parameters
----------
files : list of strings
List of paths to images.
input_shape : list
Must define what the input image's shape is.
learning_rate : float, optional
Learning rate.
batch_size : int, optional
Batch size.
n_epochs : int, optional
Number of epochs.
n_examples : int, optional
Number of example to use while demonstrating the current training
iteration's reconstruction. Creates a square montage, so make
sure int(sqrt(n_examples))**2 = n_examples, e.g. 16, 25, 36, ... 100.
crop_shape : list, optional
Size to centrally crop the image to.
crop_factor : float, optional
Resize factor to apply before cropping.
n_filters : list, optional
Same as VAE's n_filters.
n_hidden : int, optional
Same as VAE's n_hidden.
n_code : int, optional
Same as VAE's n_code.
convolutional : bool, optional
Use convolution or not.
variational : bool, optional
Use variational layer or not.
filter_sizes : list, optional
Same as VAE's filter_sizes.
dropout : bool, optional
Use dropout or not
keep_prob : float, optional
Percent of keep for dropout.
activation : function, optional
Which activation function to use.
img_step : int, optional
How often to save training images showing the manifold and
reconstruction.
save_step : int, optional
How often to save checkpoints.
ckpt_name : str, optional
Checkpoints will be named as this, e.g. 'model.ckpt'
"""
batch = create_input_pipeline(files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=crop_shape,
crop_factor=crop_factor,
shape=input_shape)
ae = VAE(input_shape=[None] + crop_shape,
convolutional=convolutional,
variational=variational,
n_filters=n_filters,
n_hidden=n_hidden,
n_code=n_code,
dropout=dropout,
filter_sizes=filter_sizes,
activation=activation)
# Create a manifold of our inner most layer to show
# example reconstructions. This is one way to see
# what the "embedding" or "latent space" of the encoder
# is capable of encoding, though note that this is just
# a random hyperplane within the latent space, and does not
# encompass all possible embeddings.
zs = np.random.uniform(-1.0, 1.0, [4, n_code]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['cost'])
# We create a session to use the graph
sess = tf.Session()
saver = tf.train.Saver()
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
# This will handle our threaded image pipeline
coord = tf.train.Coordinator()
# Ensure no more changes to graph
tf.get_default_graph().finalize()
# Start up the queues for handling the image pipeline
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
if os.path.exists(ckpt_name + '.index') or os.path.exists(ckpt_name):
saver.restore(sess, ckpt_name)
# Fit all training data
t_i = 0
batch_i = 0
epoch_i = 0
cost = 0
n_files = len(files)
test_xs = sess.run(batch) / 255.0
utils.montage(test_xs, 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
batch_i += 1
batch_xs = sess.run(batch) / 255.0
train_cost = sess.run([ae['cost'], optimizer],
feed_dict={
ae['x']: batch_xs,
ae['train']: True,
ae['keep_prob']: keep_prob
})[0]
print(batch_i, train_cost)
cost += train_cost
if batch_i % n_files == 0:
print('epoch:', epoch_i)
print('average cost:', cost / batch_i)
cost = 0
batch_i = 0
epoch_i += 1
if batch_i % img_step == 0:
# Plot example reconstructions from latent layer
recon = sess.run(ae['y'],
feed_dict={
ae['z']: zs,
ae['train']: False,
ae['keep_prob']: 1.0
})
utils.montage(recon.reshape([-1] + crop_shape),
'manifold_%08d.png' % t_i)
# Plot example reconstructions
recon = sess.run(ae['y'],
feed_dict={
ae['x']: test_xs,
ae['train']: False,
ae['keep_prob']: 1.0
})
print('reconstruction (min, max, mean):', recon.min(),
recon.max(), recon.mean())
utils.montage(recon.reshape([-1] + crop_shape),
'reconstruction_%08d.png' % t_i)
t_i += 1
if batch_i % save_step == 0:
# Save the variables to disk.
saver.save(sess,
ckpt_name,
global_step=batch_i,
write_meta_graph=False)
except tf.errors.OutOfRangeError:
print('Done.')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()
def test_mnist():
"""Train an autoencoder on MNIST.
This function will train an autoencoder on MNIST and also
save many image files during the training process, demonstrating
the latent space of the inner most dimension of the encoder,
as well as reconstructions of the decoder.
"""
# load MNIST
n_code = 2
mnist = MNIST(split=[0.8, 0.1, 0.1])
ae = VAE(input_shape=[None, 784],
n_filters=[512, 256],
n_hidden=64,
n_code=n_code,
activation=tf.nn.sigmoid,
convolutional=False,
variational=True)
n_examples = 100
zs = np.random.uniform(-1.0, 1.0, [4, n_code]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
learning_rate = 0.02
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['cost'])
# We create a session to use the graph
sess = tf.Session()
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
# Fit all training data
t_i = 0
batch_i = 0
batch_size = 200
n_epochs = 10
test_xs = mnist.test.images[:n_examples]
utils.montage(test_xs.reshape((-1, 28, 28)), 'test_xs.png')
for epoch_i in range(n_epochs):
train_i = 0
train_cost = 0
for batch_xs, _ in mnist.train.next_batch(batch_size):
train_cost += sess.run([ae['cost'], optimizer],
feed_dict={
ae['x']: batch_xs,
ae['train']: True,
ae['keep_prob']: 1.0
})[0]
train_i += 1
if batch_i % 10 == 0:
# Plot example reconstructions from latent layer
recon = sess.run(ae['y'],
feed_dict={
ae['z']: zs,
ae['train']: False,
ae['keep_prob']: 1.0
})
utils.montage(recon.reshape((-1, 28, 28)),
'manifold_%08d.png' % t_i)
# Plot example reconstructions
recon = sess.run(ae['y'],
feed_dict={
ae['x']: test_xs,
ae['train']: False,
ae['keep_prob']: 1.0
})
utils.montage(recon.reshape((-1, 28, 28)),
'reconstruction_%08d.png' % t_i)
t_i += 1
batch_i += 1
valid_i = 0
valid_cost = 0
for batch_xs, _ in mnist.valid.next_batch(batch_size):
valid_cost += sess.run([ae['cost']],
feed_dict={
ae['x']: batch_xs,
ae['train']: False,
ae['keep_prob']: 1.0
})[0]
valid_i += 1
print('train:', train_cost / train_i, 'valid:', valid_cost / valid_i)
def train_input_pipeline(
files,
A, # img_h
B, # img_w
C,
T=20,
n_enc=512,
n_z=256,
n_dec=512,
read_n=15,
write_n=15,
batch_size=64,
n_epochs=1e9,
input_shape=(64, 64, 3)):
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
batch = create_input_pipeline(
files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=(A, B, C),
shape=input_shape)
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
# Fit all training data
batch_i = 0
epoch_i = 0
n_files = len(files)
test_xs = sess.run(batch).reshape((-1, A * B * C)) / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
batch_xs = sess.run(batch) / 255.0
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={
draw['x']: batch_xs.reshape((-1, A * B * C)) / 255.0,
draw['noise']: noise
})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % n_files == 0:
batch_i = 0
epoch_i += 1
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(recon, saveto='manifold_%08d.gif' % batch_i)
plt.close('all')
batch_i += 1
except tf.errors.OutOfRangeError:
print('Done.')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()
def train_input_pipeline(files,
init_lr_g=1e-4,
init_lr_d=1e-4,
n_features=10,
n_latent=100,
n_epochs=1000000,
batch_size=200,
n_samples=15,
input_shape=[218, 178, 3],
crop_shape=[64, 64, 3],
crop_factor=0.8):
"""Summary
Parameters
----------
files : TYPE
Description
init_lr_g : float, optional
Description
init_lr_d : float, optional
Description
n_features : int, optional
Description
n_latent : int, optional
Description
n_epochs : int, optional
Description
batch_size : int, optional
Description
n_samples : int, optional
Description
input_shape : list, optional
Description
crop_shape : list, optional
Description
crop_factor : float, optional
Description
No Longer Returned
------------------
name : TYPE
Description
"""
with tf.Graph().as_default(), tf.Session() as sess:
batch = create_input_pipeline(
files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=crop_shape,
crop_factor=crop_factor,
shape=input_shape)
gan = GAN(
input_shape=[None] + crop_shape,
n_features=n_features,
n_latent=n_latent,
rgb=True,
debug=False)
vars_d = [
v for v in tf.trainable_variables()
if v.name.startswith('discriminator')
]
print('Training discriminator variables:')
[
print(v.name) for v in tf.trainable_variables()
if v.name.startswith('discriminator')
]
vars_g = [
v for v in tf.trainable_variables()
if v.name.startswith('generator')
]
print('Training generator variables:')
[
print(v.name) for v in tf.trainable_variables()
if v.name.startswith('generator')
]
zs = np.random.uniform(-1.0, 1.0, [4, n_latent]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_samples)
lr_g = tf.placeholder(tf.float32, shape=[], name='learning_rate_g')
lr_d = tf.placeholder(tf.float32, shape=[], name='learning_rate_d')
try:
from tf.contrib.layers import apply_regularization
d_reg = apply_regularization(
tf.contrib.layers.l2_regularizer(1e-6), vars_d)
g_reg = apply_regularization(
tf.contrib.layers.l2_regularizer(1e-6), vars_g)
except:
d_reg, g_reg = 0, 0
opt_g = tf.train.AdamOptimizer(
lr_g, name='Adam_g').minimize(
gan['loss_G'] + g_reg, var_list=vars_g)
opt_d = tf.train.AdamOptimizer(
lr_d, name='Adam_d').minimize(
gan['loss_D'] + d_reg, var_list=vars_d)
# %%
# We create a session to use the graph
saver = tf.train.Saver()
sums = gan['sums']
G_sum_op = tf.summary.merge([
sums['G'], sums['loss_G'], sums['z'], sums['loss_D_fake'],
sums['D_fake']
])
D_sum_op = tf.summary.merge([
sums['loss_D'], sums['loss_D_real'], sums['loss_D_fake'], sums['z'],
sums['x'], sums['D_real'], sums['D_fake']
])
writer = tf.summary.FileWriter("./logs", sess.graph_def)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
# g = tf.get_default_graph()
# [print(op.name) for op in g.get_operations()]
if os.path.exists("gan.ckpt"):
saver.restore(sess, "gan.ckpt")
print("GAN model restored.")
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
step_i, t_i = 0, 0
loss_d = 1
loss_g = 1
n_loss_d, total_loss_d = 1, 1
n_loss_g, total_loss_g = 1, 1
try:
while not coord.should_stop():
batch_xs = sess.run(batch)
step_i += 1
batch_zs = np.random.uniform(
-1.0, 1.0, [batch_size, n_latent]).astype(np.float32)
this_lr_g = min(1e-2,
max(1e-6, init_lr_g * (loss_g / loss_d)**2))
this_lr_d = min(1e-2,
max(1e-6, init_lr_d * (loss_d / loss_g)**2))
# this_lr_d *= ((1.0 - (step_i / 100000)) ** 2)
# this_lr_g *= ((1.0 - (step_i / 100000)) ** 2)
# if np.random.random() > (loss_g / (loss_d + loss_g)):
if step_i % 3 == 1:
loss_d, _, sum_d = sess.run(
[gan['loss_D'], opt_d, D_sum_op],
feed_dict={
gan['x']: batch_xs,
gan['z']: batch_zs,
lr_d: this_lr_d
})
total_loss_d += loss_d
n_loss_d += 1
writer.add_summary(sum_d, step_i)
print('%04d d* = lr: %0.08f, loss: %08.06f, \t' % (
step_i, this_lr_d, loss_d
) + 'g = lr: %0.08f, loss: %08.06f' % (this_lr_g, loss_g))
else:
loss_g, _, sum_g = sess.run(
[gan['loss_G'], opt_g, G_sum_op],
feed_dict={gan['z']: batch_zs,
lr_g: this_lr_g})
total_loss_g += loss_g
n_loss_g += 1
writer.add_summary(sum_g, step_i)
print('%04d d = lr: %0.08f, loss: %08.06f, \t' % (
step_i, this_lr_d, loss_d
) + 'g* = lr: %0.08f, loss: %08.06f' % (this_lr_g, loss_g))
if step_i % 100 == 0:
samples = sess.run(gan['G'], feed_dict={gan['z']: zs})
utils.montage(
np.clip((samples + 1) * 127.5, 0, 255).astype(np.uint8),
'imgs/gan_%08d.png' % t_i)
t_i += 1
print('generator loss:', total_loss_g / n_loss_g)
print('discriminator loss:', total_loss_d / n_loss_d)
# Save the variables to disk.
save_path = saver.save(
sess,
"./gan.ckpt",
global_step=step_i,
write_meta_graph=False)
print("Model saved in file: %s" % save_path)
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
def test_mnist():
A = 28 # img_h
B = 28 # img_w
C = 1
T = 10
n_enc = 256
n_z = 100
n_dec = 256
read_n = 5
write_n = 5
batch_size = 64
mnist = MNIST(split=[0.8, 0.1, 0.1])
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
n_epochs = 100
test_xs = mnist.test.images[:n_examples]
utils.montage(test_xs.reshape((-1, A, B)), 'test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in mnist.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={draw['x']: batch_xs,
draw['noise']: noise})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [utils.montage(r.reshape(-1, A, B)) for r in recon]
gif.build_gif(
recon,
cmap='gray',
saveto='manifold_%08d.gif' % batch_i)
saver.save(sess, './draw.ckpt', global_step=batch_i)
batch_i += 1
def train_dataset(ds,
A,
B,
C,
T=20,
n_enc=512,
n_z=200,
n_dec=512,
read_n=12,
write_n=12,
batch_size=100,
n_epochs=100):
if ds is None:
ds = CIFAR10(split=[0.8, 0.1, 0.1])
A, B, C = (32, 32, 3)
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
# Clip gradients
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
# Add summary variables
tf.summary.scalar(name='cost', tensor=draw['cost'])
tf.summary.scalar(name='loss_z', tensor=draw['loss_z'])
tf.summary.scalar(name='loss_x', tensor=draw['loss_x'])
tf.summary.histogram(
name='recon_t0_histogram', values=draw['canvas'][0])
tf.summary.histogram(
name='recon_t-1_histogram', values=draw['canvas'][-1])
tf.summary.image(
name='recon_t0_image',
tensor=tf.reshape(draw['canvas'][0], (-1, A, B, C)),
max_outputs=2)
tf.summary.image(
name='recon_t-1_image',
tensor=tf.reshape(draw['canvas'][-1], (-1, A, B, C)),
max_outputs=2)
sums = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(logdir='draw/train')
valid_writer = tf.summary.FileWriter(logdir='draw/valid')
# Init
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
test_xs = ds.test.images[:n_examples] / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'draw/test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in ds.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums, train_op],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
train_writer.add_summary(summary, batch_i)
print('train cost:', cost)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(
recon,
cmap='gray',
saveto='draw/manifold_%08d.gif' % batch_i)
saver.save(sess, './draw/draw.ckpt', global_step=batch_i)
batch_i += 1
# Run validation
if batch_i % 1000 == 0:
for batch_xs, _ in ds.valid.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
valid_writer.add_summary(summary, batch_i)
print('valid cost:', cost)
batch_i += 1
def train_input_pipeline(files,
init_lr_g=1e-4,
init_lr_d=1e-4,
n_features=10,
n_latent=100,
n_epochs=1000000,
batch_size=200,
n_samples=15,
input_shape=[218, 178, 3],
crop_shape=[64, 64, 3],
crop_factor=0.8):
"""Summary
Parameters
----------
files : TYPE
Description
init_lr_g : float, optional
Description
init_lr_d : float, optional
Description
n_features : int, optional
Description
n_latent : int, optional
Description
n_epochs : int, optional
Description
batch_size : int, optional
Description
n_samples : int, optional
Description
input_shape : list, optional
Description
crop_shape : list, optional
Description
crop_factor : float, optional
Description
No Longer Returned
------------------
name : TYPE
Description
"""
with tf.Graph().as_default(), tf.Session() as sess:
batch = create_input_pipeline(files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=crop_shape,
crop_factor=crop_factor,
shape=input_shape)
gan = GAN(input_shape=[None] + crop_shape,
n_features=n_features,
n_latent=n_latent,
rgb=True,
debug=False)
vars_d = [
v for v in tf.trainable_variables()
if v.name.startswith('discriminator')
]
print('Training discriminator variables:')
[
print(v.name) for v in tf.trainable_variables()
if v.name.startswith('discriminator')
]
vars_g = [
v for v in tf.trainable_variables()
if v.name.startswith('generator')
]
print('Training generator variables:')
[
print(v.name) for v in tf.trainable_variables()
if v.name.startswith('generator')
]
zs = np.random.uniform(-1.0, 1.0, [4, n_latent]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_samples)
lr_g = tf.placeholder(tf.float32, shape=[], name='learning_rate_g')
lr_d = tf.placeholder(tf.float32, shape=[], name='learning_rate_d')
try:
from tf.contrib.layers import apply_regularization
d_reg = apply_regularization(
tf.contrib.layers.l2_regularizer(1e-6), vars_d)
g_reg = apply_regularization(
tf.contrib.layers.l2_regularizer(1e-6), vars_g)
except:
d_reg, g_reg = 0, 0
opt_g = tf.train.AdamOptimizer(lr_g, name='Adam_g').minimize(
gan['loss_G'] + g_reg, var_list=vars_g)
opt_d = tf.train.AdamOptimizer(lr_d, name='Adam_d').minimize(
gan['loss_D'] + d_reg, var_list=vars_d)
# %%
# We create a session to use the graph
saver = tf.train.Saver()
sums = gan['sums']
G_sum_op = tf.summary.merge([
sums['G'], sums['loss_G'], sums['z'], sums['loss_D_fake'],
sums['D_fake']
])
D_sum_op = tf.summary.merge([
sums['loss_D'], sums['loss_D_real'], sums['loss_D_fake'],
sums['z'], sums['x'], sums['D_real'], sums['D_fake']
])
writer = tf.summary.FileWriter("./logs", sess.graph_def)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
# g = tf.get_default_graph()
# [print(op.name) for op in g.get_operations()]
if os.path.exists("gan.ckpt"):
saver.restore(sess, "gan.ckpt")
print("GAN model restored.")
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
step_i, t_i = 0, 0
loss_d = 1
loss_g = 1
n_loss_d, total_loss_d = 1, 1
n_loss_g, total_loss_g = 1, 1
try:
while not coord.should_stop():
batch_xs = sess.run(batch)
step_i += 1
batch_zs = np.random.uniform(
-1.0, 1.0, [batch_size, n_latent]).astype(np.float32)
this_lr_g = min(1e-2,
max(1e-6,
init_lr_g * (loss_g / loss_d)**2))
this_lr_d = min(1e-2,
max(1e-6,
init_lr_d * (loss_d / loss_g)**2))
# this_lr_d *= ((1.0 - (step_i / 100000)) ** 2)
# this_lr_g *= ((1.0 - (step_i / 100000)) ** 2)
# if np.random.random() > (loss_g / (loss_d + loss_g)):
if step_i % 3 == 1:
loss_d, _, sum_d = sess.run(
[gan['loss_D'], opt_d, D_sum_op],
feed_dict={
gan['x']: batch_xs,
gan['z']: batch_zs,
lr_d: this_lr_d
})
total_loss_d += loss_d
n_loss_d += 1
writer.add_summary(sum_d, step_i)
print('%04d d* = lr: %0.08f, loss: %08.06f, \t' %
(step_i, this_lr_d, loss_d) +
'g = lr: %0.08f, loss: %08.06f' %
(this_lr_g, loss_g))
else:
loss_g, _, sum_g = sess.run(
[gan['loss_G'], opt_g, G_sum_op],
feed_dict={
gan['z']: batch_zs,
lr_g: this_lr_g
})
total_loss_g += loss_g
n_loss_g += 1
writer.add_summary(sum_g, step_i)
print('%04d d = lr: %0.08f, loss: %08.06f, \t' %
(step_i, this_lr_d, loss_d) +
'g* = lr: %0.08f, loss: %08.06f' %
(this_lr_g, loss_g))
if step_i % 100 == 0:
samples = sess.run(gan['G'], feed_dict={gan['z']: zs})
utils.montage(
np.clip((samples + 1) * 127.5, 0,
255).astype(np.uint8),
'imgs/gan_%08d.png' % t_i)
t_i += 1
print('generator loss:', total_loss_g / n_loss_g)
print('discriminator loss:', total_loss_d / n_loss_d)
# Save the variables to disk.
save_path = saver.save(sess,
"./gan.ckpt",
global_step=step_i,
write_meta_graph=False)
print("Model saved in file: %s" % save_path)
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
def train_input_pipeline(
files,
A, # img_h
B, # img_w
C,
T=20,
n_enc=512,
n_z=256,
n_dec=512,
read_n=15,
write_n=15,
batch_size=64,
n_epochs=1e9,
input_shape=(64, 64, 3)):
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
batch = create_input_pipeline(
files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=(A, B, C),
shape=input_shape)
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
# Fit all training data
batch_i = 0
epoch_i = 0
n_files = len(files)
test_xs = sess.run(batch).reshape((-1, A * B * C)) / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
batch_xs = sess.run(batch) / 255.0
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={
draw['x']: batch_xs.reshape((-1, A * B * C)) / 255.0,
draw['noise']: noise
})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % n_files == 0:
batch_i = 0
epoch_i += 1
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(recon, saveto='manifold_%08d.gif' % batch_i)
plt.close('all')
batch_i += 1
except tf.errors.OutOfRangeError:
print('Done.')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()
def train_dataset(ds,
A,
B,
C,
T=20,
n_enc=512,
n_z=200,
n_dec=512,
read_n=12,
write_n=12,
batch_size=100,
n_epochs=100):
if ds is None:
ds = CIFAR10(split=[0.8, 0.1, 0.1])
A, B, C = (32, 32, 3)
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
# Clip gradients
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
# Add summary variables
tf.summary.scalar(name='cost', tensor=draw['cost'])
tf.summary.scalar(name='loss_z', tensor=draw['loss_z'])
tf.summary.scalar(name='loss_x', tensor=draw['loss_x'])
tf.summary.histogram(
name='recon_t0_histogram', values=draw['canvas'][0])
tf.summary.histogram(
name='recon_t-1_histogram', values=draw['canvas'][-1])
tf.summary.image(
name='recon_t0_image',
tensor=tf.reshape(draw['canvas'][0], (-1, A, B, C)),
max_outputs=2)
tf.summary.image(
name='recon_t-1_image',
tensor=tf.reshape(draw['canvas'][-1], (-1, A, B, C)),
max_outputs=2)
sums = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(logdir='draw/train')
valid_writer = tf.summary.FileWriter(logdir='draw/valid')
# Init
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
test_xs = ds.test.images[:n_examples] / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'draw/test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in ds.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums, train_op],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
train_writer.add_summary(summary, batch_i)
print('train cost:', cost)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(
recon,
cmap='gray',
saveto='draw/manifold_%08d.gif' % batch_i)
saver.save(sess, './draw/draw.ckpt', global_step=batch_i)
batch_i += 1
# Run validation
if batch_i % 1000 == 0:
for batch_xs, _ in ds.valid.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
valid_writer.add_summary(summary, batch_i)
print('valid cost:', cost)
batch_i += 1
def test_mnist():
A = 28 # img_h
B = 28 # img_w
C = 1
T = 10
n_enc = 256
n_z = 100
n_dec = 256
read_n = 5
write_n = 5
batch_size = 64
mnist = MNIST(split=[0.8, 0.1, 0.1])
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
n_epochs = 100
test_xs = mnist.test.images[:n_examples]
utils.montage(test_xs.reshape((-1, A, B)), 'test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in mnist.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={draw['x']: batch_xs,
draw['noise']: noise})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [utils.montage(r.reshape(-1, A, B)) for r in recon]
gif.build_gif(
recon,
cmap='gray',
saveto='manifold_%08d.gif' % batch_i)
saver.save(sess, './draw.ckpt', global_step=batch_i)
batch_i += 1
def train_vaegan(files,
learning_rate=0.00001,
batch_size=64,
n_epochs=250,
n_examples=10,
input_shape=[218, 178, 3],
crop_shape=[64, 64, 3],
crop_factor=0.8,
n_filters=[100, 100, 100, 100],
n_hidden=None,
n_code=128,
convolutional=True,
variational=True,
filter_sizes=[3, 3, 3, 3],
activation=tf.nn.elu,
ckpt_name="vaegan.ckpt"):
"""Summary
Parameters
----------
files : TYPE
Description
learning_rate : float, optional
Description
batch_size : int, optional
Description
n_epochs : int, optional
Description
n_examples : int, optional
Description
input_shape : list, optional
Description
crop_shape : list, optional
Description
crop_factor : float, optional
Description
n_filters : list, optional
Description
n_hidden : int, optional
Description
n_code : int, optional
Description
convolutional : bool, optional
Description
variational : bool, optional
Description
filter_sizes : list, optional
Description
activation : TYPE, optional
Description
ckpt_name : str, optional
Description
No Longer Returned
------------------
name : TYPE
Description
"""
ae = VAEGAN(
input_shape=[None] + crop_shape,
convolutional=convolutional,
variational=variational,
n_filters=n_filters,
n_hidden=n_hidden,
n_code=n_code,
filter_sizes=filter_sizes,
activation=activation)
batch = create_input_pipeline(
files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=crop_shape,
crop_factor=crop_factor,
shape=input_shape)
zs = np.random.randn(4, n_code).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
opt_enc = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['loss_enc'],
var_list=[
var_i for var_i in tf.trainable_variables()
if var_i.name.startswith('encoder')
])
opt_gen = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['loss_gen'],
var_list=[
var_i for var_i in tf.trainable_variables()
if var_i.name.startswith('generator')
])
opt_dis = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
ae['loss_dis'],
var_list=[
var_i for var_i in tf.trainable_variables()
if var_i.name.startswith('discriminator')
])
sess = tf.Session()
saver = tf.train.Saver()
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
if os.path.exists(ckpt_name + '.index') or os.path.exists(ckpt_name):
saver.restore(sess, ckpt_name)
print("VAE model restored.")
t_i = 0
batch_i = 0
epoch_i = 0
equilibrium = 0.693
margin = 0.4
n_files = len(files)
test_xs = sess.run(batch) / 255.0
utils.montage(test_xs, 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
if batch_i % (n_files // batch_size) == 0:
batch_i = 0
epoch_i += 1
print('---------- EPOCH:', epoch_i)
batch_i += 1
batch_xs = sess.run(batch) / 255.0
batch_zs = np.random.randn(batch_size, n_code).astype(np.float32)
real_cost, fake_cost, _ = sess.run(
[ae['loss_real'], ae['loss_fake'], opt_enc],
feed_dict={ae['x']: batch_xs,
ae['gamma']: 0.5})
real_cost = -np.mean(real_cost)
fake_cost = -np.mean(fake_cost)
print('real:', real_cost, '/ fake:', fake_cost)
gen_update = True
dis_update = True
if real_cost > (equilibrium + margin) or \
fake_cost > (equilibrium + margin):
gen_update = False
if real_cost < (equilibrium - margin) or \
fake_cost < (equilibrium - margin):
dis_update = False
if not (gen_update or dis_update):
gen_update = True
dis_update = True
if gen_update:
sess.run(
opt_gen,
feed_dict={
ae['x']: batch_xs,
ae['z_samp']: batch_zs,
ae['gamma']: 0.5
})
if dis_update:
sess.run(
opt_dis,
feed_dict={
ae['x']: batch_xs,
ae['z_samp']: batch_zs,
ae['gamma']: 0.5
})
if batch_i % 50 == 0:
# Plot example reconstructions from latent layer
recon = sess.run(ae['x_tilde'], feed_dict={ae['z']: zs})
print('recon:', recon.min(), recon.max())
recon = np.clip(recon / recon.max(), 0, 1)
utils.montage(
recon.reshape([-1] + crop_shape),
'imgs/manifold_%08d.png' % t_i)
# Plot example reconstructions
recon = sess.run(ae['x_tilde'], feed_dict={ae['x']: test_xs})
print('recon:', recon.min(), recon.max())
recon = np.clip(recon / recon.max(), 0, 1)
utils.montage(
recon.reshape([-1] + crop_shape),
'imgs/reconstruction_%08d.png' % t_i)
t_i += 1
if batch_i % 100 == 0:
# Save the variables to disk.
save_path = saver.save(
sess,
ckpt_name,
global_step=batch_i,
write_meta_graph=False)
print("Model saved in file: %s" % save_path)
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()