def train(options):
# Placeholders for input data and the targets
with tf.name_scope('Input'):
X = tf.placeholder(dtype=tf.float32,
shape=[options.batch_size, input_dim],
name='Input')
X_noisy = tf.placeholder(dtype=tf.float32,
shape=[options.batch_size, input_dim],
name='Input_noisy')
input_images = tf.reshape(X_noisy, [-1, 28, 28, 1])
with tf.name_scope('Latent_variable'):
z = tf.placeholder(dtype=tf.float32,
shape=[None, options.z_dim],
name='Latent_variable')
with tf.variable_scope('Encoder'):
encoder_output = encoder(X_noisy)
with tf.variable_scope('Decoder') as scope:
decoder_output, logits = decoder(encoder_output)
generated_images = tf.reshape(decoder_output, [-1, 28, 28, 1])
scope.reuse_variables()
X_samples, _ = decoder(z)
with tf.name_scope('Loss'):
loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=X))
# Optimizer
train_op, grads_and_vars = AdamOptimizer(loss, options.learning_rate,
options.beta1)
# Visualization
tf.summary.scalar(name='Loss', tensor=loss)
tf.summary.histogram(name='Latent_variable', values=encoder_output)
for grad, var in grads_and_vars:
tf.summary.histogram('Gradients/' + var.name, grad)
tf.summary.histogram('Values/' + var.name, var)
tf.summary.image(name='Input Images', tensor=input_images, max_outputs=10)
tf.summary.image(name='Generated Images',
tensor=generated_images,
max_outputs=10)
summary_op = tf.summary.merge_all()
saver = tf.train.Saver()
step = 0
init = tf.global_variables_initializer()
n_batches = int(mnist.train.num_examples / options.batch_size)
with tf.Session() as sess:
sess.run(init)
if not options.run_inference:
try:
writer = tf.summary.FileWriter(logdir=options.tensorboard_path,
graph=sess.graph)
for epoch in range(options.epochs):
for iteration in range(n_batches):
batch_x, _ = mnist.train.next_batch(options.batch_size)
mask = generate_mask(batch_x.shape, noise_length)
noisy_batch = batch_x * mask
# Train
sess.run(train_op,
feed_dict={
X: batch_x,
X_noisy: noisy_batch
})
if iteration % 50 == 0:
summary, batch_loss = sess.run([summary_op, loss],
feed_dict={
X: batch_x,
X_noisy:
noisy_batch
})
writer.add_summary(summary, global_step=step)
print("Epoch: {} - Iteration {} - Loss: {:.4f}\n".
format(epoch, iteration, batch_loss))
step += 1
saver.save(sess,
save_path=options.checkpoints_path,
global_step=step)
print("Model Trained!")
except KeyboardInterrupt:
print('Stopping training...')
saver.save(sess,
save_path=options.checkpoints_path,
global_step=step)
else:
print('Restoring latest saved TensorFlow model...')
dir_path = os.path.dirname(os.path.realpath(__file__))
cur_dir = dir_path.split('/')[-1]
script_name = os.path.basename(__file__).split('.')[0]
experiments = glob.glob(
os.path.join(options.MAIN_PATH, cur_dir) +
'/{}*'.format(script_name))
experiments.sort(key=lambda x: os.path.getmtime(x))
if len(experiments) > 0:
print('Restoring: {}'.format(experiments[-1]))
saver.restore(
sess,
tf.train.latest_checkpoint(
os.path.join(experiments[-1], 'checkpoints')))
plot(sess, z, X_samples, num_images=15, height=28, width=28)
else:
print('No checkpoint found at {}'.format(
os.path.join(options.MAIN_PATH, cur_dir)))
def train(options):
# Placeholders for input data and the targets
with tf.name_scope('Input'):
X = tf.placeholder(dtype=tf.float32,
shape=[options.batch_size, input_dim],
name='Input')
input_images = tf.reshape(X, [-1, 28, 28, 1])
with tf.name_scope('Latent_variable'):
z = tf.placeholder(dtype=tf.float32,
shape=[None, options.z_dim],
name='Latent_variable')
with tf.variable_scope('Encoder'):
encoder_output = encoder(X)
with tf.variable_scope('Decoder') as scope:
decoder_prob, decoder_logits = decoder(encoder_output)
generated_images = tf.reshape(decoder_prob, [-1, 28, 28, 1])
scope.reuse_variables()
X_samples, _ = decoder(z)
with tf.variable_scope('Discriminator') as scope:
D_real = discriminator(X)
scope.reuse_variables()
D_fake = discriminator(decoder_prob)
# Loss - E[log P(X|z)]
with tf.name_scope('Loss'):
with tf.name_scope('Reconstruction_Loss'):
loss_reconstr = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(logits=decoder_logits,
labels=X))
with tf.name_scope('Discriminator_loss'):
loss_discriminator = (tf.reduce_mean(D_fake) -
tf.reduce_mean(D_real))
with tf.name_scope('Autoencoder_loss'):
loss_autoencoder = -tf.reduce_mean(
D_fake) + lambda_reconstruction * loss_reconstr
with tf.variable_scope('Discriminator_Accuracy'):
accuracy_real = tf.reduce_mean(
tf.cast(tf.greater_equal(D_real, 0.5), tf.float16))
accuracy_fake = tf.reduce_mean(
tf.cast(tf.less(D_fake, 0.5), tf.float16))
accuracy_tot = (accuracy_real + accuracy_fake) / 2
vars = tf.trainable_variables()
enc_params = [v for v in vars if v.name.startswith('Encoder/')]
dec_params = [v for v in vars if v.name.startswith('Decoder/')]
dis_params = [v for v in vars if v.name.startswith('Discriminator/')]
dis_weights = [w for w in dis_params if 'weight' in w.name]
clipped_weights = clip_weights(dis_weights, clipping_parameter,
'clip_weights')
with tf.name_scope('Optimizer'):
train_op_Dis, Dis_grads_and_vars = AdamOptimizer(loss_discriminator,
options.learning_rate,
options.beta1,
var_list=dis_params)
train_op_autoenc, Autoenc_grads_and_vars = AdamOptimizer(
loss_autoencoder,
options.learning_rate,
options.beta1,
var_list=enc_params + dec_params)
# Visualization
tf.summary.scalar(name='Discriminator_loss', tensor=loss_discriminator)
tf.summary.scalar(name='Autoencoder_loss', tensor=loss_autoencoder)
tf.summary.scalar(name='Discriminator_Accuracy', tensor=accuracy_tot)
tf.summary.histogram(name='Latent_variable', values=encoder_output)
for grad, var in Autoenc_grads_and_vars + Dis_grads_and_vars:
tf.summary.histogram('Gradients/' + var.name, grad)
tf.summary.histogram('Values/' + var.name, var)
tf.summary.image(name='Input Images', tensor=input_images, max_outputs=10)
tf.summary.image(name='Generated Images',
tensor=generated_images,
max_outputs=10)
summary_op = tf.summary.merge_all()
saver = tf.train.Saver()
step = 0
init = tf.global_variables_initializer()
n_batches = int(mnist.train.num_examples / options.batch_size)
with tf.Session() as sess:
sess.run(init)
if not options.run_inference:
try:
writer = tf.summary.FileWriter(logdir=options.tensorboard_path,
graph=sess.graph)
for epoch in range(options.epochs):
for iteration in range(n_batches):
# update discriminator
for _ in range(n_critic):
batch_x, _ = mnist.train.next_batch(
options.batch_size)
sess.run(train_op_Dis, feed_dict={X: batch_x})
_ = sess.run(clipped_weights)
# update autoencoder/generator parameters
batch_x, _ = mnist.train.next_batch(options.batch_size)
sess.run(train_op_autoenc, feed_dict={X: batch_x})
if iteration % 50 == 0:
summary, ld, la, acc = sess.run(
[
summary_op, loss_discriminator,
loss_autoencoder, accuracy_tot
],
feed_dict={X: batch_x})
writer.add_summary(summary, global_step=step)
print(
"Epoch: {} - Iteration {} - Autoencoder loss: {:.4f}"
.format(epoch, iteration, la))
print(
"Discriminator loss: {} - Discriminator accuracy: {:.4f}%\n"
.format(ld, acc * 100))
step += 1
saver.save(sess,
save_path=options.checkpoints_path,
global_step=step)
print("Model Trained!")
except KeyboardInterrupt:
print('Stopping training...')
saver.save(sess,
save_path=options.checkpoints_path,
global_step=step)
else:
print('Restoring latest saved TensorFlow model...')
dir_path = os.path.dirname(os.path.realpath(__file__))
cur_dir = dir_path.split('/')[-1]
script_name = os.path.basename(__file__).split('.')[0]
experiments = glob.glob(
os.path.join(options.MAIN_PATH, cur_dir) +
'/{}*'.format(script_name))
experiments.sort(key=lambda x: os.path.getmtime(x))
if len(experiments) > 0:
print('Restoring: {}'.format(experiments[-1]))
saver.restore(
sess,
tf.train.latest_checkpoint(
os.path.join(experiments[-1], 'checkpoints')))
plot(sess, z, X_samples, num_images=15, height=28, width=28)
else:
print('No checkpoint found at {}'.format(
os.path.join(options.MAIN_PATH, cur_dir)))
def train(options):
with tf.name_scope('Input'):
X = tf.placeholder(tf.float32, shape=[None, input_dim])
input_images = tf.reshape(X, [-1, 28, 28, 1])
with tf.name_scope('Latent_variable'):
z = tf.placeholder(tf.float32, shape=[None, options.z_dim])
with tf.variable_scope('Encoder'):
z_mu, z_logvar = encoder(X)
with tf.variable_scope('Decoder') as scope:
z_sample = sample_z(z_mu, z_logvar)
decoder_output, logits = decoder(z_sample)
generated_images = tf.reshape(decoder_output, [-1, 28, 28, 1])
scope.reuse_variables()
# Sampling from random z
X_samples, _ = decoder(z)
with tf.name_scope('Loss'):
# E[log P(X|z)]
reconstruction_loss = tf.reduce_sum(
tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=X),
1)
# D_KL(Q(z|X) || P(z|X)); calculate in closed form as both dist. are Gaussian
kl_loss = 0.5 * tf.reduce_sum(
tf.exp(z_logvar) + z_mu**2 - 1. - z_logvar, 1)
# VAE loss
vae_loss = tf.reduce_mean(reconstruction_loss + kl_loss)
# Optimizer
train_op, grads_and_vars = AdamOptimizer(vae_loss, options.learning_rate,
options.beta1)
# Visualization
tf.summary.scalar(name='Loss', tensor=vae_loss)
tf.summary.histogram(name='Sampled variable', values=z_sample)
for grad, var in grads_and_vars:
tf.summary.histogram('Gradients/' + var.name, grad)
tf.summary.histogram('Values/' + var.name, var)
tf.summary.image(name='Input Images', tensor=input_images, max_outputs=10)
tf.summary.image(name='Generated Images',
tensor=generated_images,
max_outputs=10)
summary_op = tf.summary.merge_all()
saver = tf.train.Saver()
step = 0
init = tf.global_variables_initializer()
n_batches = int(mnist.train.num_examples / options.batch_size)
with tf.Session() as sess:
sess.run(init)
if not options.run_inference:
try:
writer = tf.summary.FileWriter(logdir=options.tensorboard_path,
graph=sess.graph)
for epoch in range(options.epochs):
for iteration in range(n_batches):
batch_x, _ = mnist.train.next_batch(options.batch_size)
# Train
sess.run(train_op, feed_dict={X: batch_x})
if iteration % 50 == 0:
summary, batch_loss = sess.run(
[summary_op, vae_loss], feed_dict={X: batch_x})
writer.add_summary(summary, global_step=step)
print("Epoch: {} - Iteration {} - Loss: {:.4f}\n".
format(epoch, iteration, batch_loss))
step += 1
saver.save(sess,
save_path=options.checkpoints_path,
global_step=step)
print("Model Trained!")
except KeyboardInterrupt:
print('Stopping training...')
saver.save(sess,
save_path=options.checkpoints_path,
global_step=step)
else:
print('Restoring latest saved TensorFlow model...')
dir_path = os.path.dirname(os.path.realpath(__file__))
cur_dir = dir_path.split('/')[-1]
script_name = os.path.basename(__file__).split('.')[0]
experiments = glob.glob(
os.path.join(options.MAIN_PATH, cur_dir) +
'/{}*'.format(script_name))
experiments.sort(key=lambda x: os.path.getmtime(x))
if len(experiments) > 0:
print('Restoring: {}'.format(experiments[-1]))
saver.restore(
sess,
tf.train.latest_checkpoint(
os.path.join(experiments[-1], 'checkpoints')))
plot(sess, z, X_samples, num_images=15, height=28, width=28)
else:
print('No checkpoint found at {}'.format(
os.path.join(options.MAIN_PATH, cur_dir)))
