def main(hparams):
# Set up some stuff according to hparams
utils.set_up_dir(hparams.ckpt_dir)
utils.set_up_dir(hparams.sample_dir)
utils.print_hparams(hparams)
# encode
x_ph = tf.placeholder(tf.float32, [None, hparams.n_input], name='x_ph')
z_mean, z_log_sigma_sq = model_def.encoder(hparams,
x_ph,
'enc',
reuse=False)
# sample
eps = tf.random_normal((hparams.batch_size, hparams.n_z),
0,
1,
dtype=tf.float32)
z_sigma = tf.sqrt(tf.exp(z_log_sigma_sq))
z = z_mean + z_sigma * eps
# reconstruct
logits, x_reconstr_mean = model_def.generator(hparams,
z,
'gen',
reuse=False)
# generator sampler
z_ph = tf.placeholder(tf.float32, [None, hparams.n_z], name='x_ph')
_, x_sample = model_def.generator(hparams, z_ph, 'gen', reuse=True)
# define loss and update op
total_loss = model_def.get_loss(x_ph, logits, z_mean, z_log_sigma_sq)
opt = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
update_op = opt.minimize(total_loss)
# Sanity checks
for var in tf.global_variables():
print var.op.name
print ''
# Get a new session
sess = tf.Session()
# Model checkpointing setup
model_saver = tf.train.Saver()
init_op = tf.global_variables_initializer()
sess.run(init_op)
# Attempt to restore variables from checkpoint
start_epoch = utils.try_restore(hparams, sess, model_saver)
# Get data iterator
iterator = data_input.channel_data_iteratior(hparams)
next_element = iterator.get_next()
# Training
for epoch in range(start_epoch + 1, hparams.training_epochs):
avg_loss = 0.0
num_batches = hparams.num_samples // hparams.batch_size
batch_num = 0
for i in range(num_batches):
try:
x_batch_val = sess.run(next_element['H_data'])
x_batch_val = np.squeeze(x_batch_val)
batch_num += 1
feed_dict = {x_ph: x_batch_val}
_, loss_val = sess.run([update_op, total_loss],
feed_dict=feed_dict)
#print(loss_val)
avg_loss += loss_val / hparams.num_samples * hparams.batch_size
except tf.errors.OutOfRangeError:
print("End of dataset")
break
if epoch % hparams.summary_epoch == 0:
print "Epoch:", '%04d' % (epoch), 'Avg loss = {:.9f}'.format(
avg_loss)
if epoch % hparams.ckpt_epoch == 0:
save_path = os.path.join(hparams.ckpt_dir, 'channel_vae_model')
model_saver.save(sess, save_path, global_step=epoch)
save_path = os.path.join(hparams.ckpt_dir, 'channel_vae_model')
model_saver.save(sess, save_path, global_step=hparams.training_epochs - 1)
def main(hparams):
# Set up some stuff according to hparams
utils.set_up_dir(hparams.ckpt_dir)
utils.set_up_dir(hparams.sample_dir)
utils.print_hparams(hparams)
# encode
x_ph = tf.placeholder(tf.float32, [None, hparams.n_input], name='x_ph')
z_mean, z_log_sigma_sq = model_def.encoder(hparams,
x_ph,
'enc',
reuse=False)
# sample
eps = tf.random_normal((hparams.batch_size, hparams.n_z),
0,
1,
dtype=tf.float32)
z_sigma = tf.sqrt(tf.exp(z_log_sigma_sq))
z = z_mean + z_sigma * eps
# reconstruct
logits, x_reconstr_mean = model_def.generator(hparams,
z,
'gen',
reuse=False)
# generator sampler
z_ph = tf.placeholder(tf.float32, [None, hparams.n_z], name='x_ph')
_, x_sample = model_def.generator(hparams, z_ph, 'gen', reuse=True)
# define loss and update op
total_loss = model_def.get_loss(x_ph, logits, z_mean, z_log_sigma_sq)
opt = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
update_op = opt.minimize(total_loss)
# Sanity checks
for var in tf.global_variables():
print var.op.name
print ''
# Get a new session
sess = tf.Session()
# Model checkpointing setup
model_saver = tf.train.Saver()
init_op = tf.global_variables_initializer()
sess.run(init_op)
# Attempt to restore variables from checkpoint
start_epoch = utils.try_restore(hparams, sess, model_saver)
# Get data iterator
iterator = data_input.omniglot_data_iterator()
# Training
for epoch in range(start_epoch + 1, hparams.training_epochs):
avg_loss = 0.0
num_batches = hparams.num_samples // hparams.batch_size
batch_num = 0
for (x_batch_val, _) in iterator(hparams, num_batches):
batch_num += 1
feed_dict = {x_ph: x_batch_val}
_, loss_val = sess.run([update_op, total_loss],
feed_dict=feed_dict)
avg_loss += loss_val / hparams.num_samples * hparams.batch_size
if batch_num % 100 == 0:
x_reconstr_mean_val = sess.run(x_reconstr_mean,
feed_dict={x_ph: x_batch_val})
z_val = np.random.randn(hparams.batch_size, hparams.n_z)
x_sample_val = sess.run(x_sample, feed_dict={z_ph: z_val})
utils.save_images(
np.reshape(x_reconstr_mean_val, [-1, 28, 28]), [10, 10],
'{}/reconstr_{:02d}_{:04d}.png'.format(
hparams.sample_dir, epoch, batch_num))
utils.save_images(
np.reshape(x_batch_val, [-1, 28, 28]), [10, 10],
'{}/orig_{:02d}_{:04d}.png'.format(hparams.sample_dir,
epoch, batch_num))
utils.save_images(
np.reshape(x_sample_val, [-1, 28, 28]), [10, 10],
'{}/sampled_{:02d}_{:04d}.png'.format(
hparams.sample_dir, epoch, batch_num))
if epoch % hparams.summary_epoch == 0:
print "Epoch:", '%04d' % (epoch), 'Avg loss = {:.9f}'.format(
avg_loss)
if epoch % hparams.ckpt_epoch == 0:
save_path = os.path.join(hparams.ckpt_dir, 'omniglot_vae_model')
model_saver.save(sess, save_path, global_step=epoch)
save_path = os.path.join(hparams.ckpt_dir, 'omniglot_vae_model')
model_saver.save(sess, save_path, global_step=hparams.training_epochs - 1)