def build(self):
"""Build the TF graph and heads for dataspace model.
It also prepares different graph, session and heads for sampling and
classification respectively.
"""
config_name = self.config_name
config = common.load_config(config_name)
exp_uid = self.exp_uid
graph = tf.Graph()
with graph.as_default():
sess = tf.Session(graph=graph)
config = common.load_config(config_name)
model_uid = common.get_model_uid(config_name, exp_uid)
m = model_dataspace.Model(config, name=model_uid)
m()
self.config = config
self.graph = graph
self.sess = sess
self.m = m
def main(unused_argv):
del unused_argv
# Load Config
config_name = FLAGS.config
config_module = importlib.import_module(configs_module_prefix +
'.%s' % config_name)
config = config_module.config
model_uid = common.get_model_uid(config_name, FLAGS.exp_uid)
batch_size = config['batch_size']
# Load dataset
dataset = common.load_dataset(config)
save_path = dataset.save_path
train_data = dataset.train_data
attr_train = dataset.attr_train
eval_data = dataset.eval_data
attr_eval = dataset.attr_eval
# Make the directory
save_dir = os.path.join(save_path, model_uid)
best_dir = os.path.join(save_dir, 'best')
tf.gfile.MakeDirs(save_dir)
tf.gfile.MakeDirs(best_dir)
tf.logging.info('Save Dir: %s', save_dir)
np.random.seed(FLAGS.random_seed)
# We use `N` in variable name to emphasis its being the Number of something.
N_train = train_data.shape[0] # pylint:disable=invalid-name
N_eval = eval_data.shape[0] # pylint:disable=invalid-name
# Load Model
tf.reset_default_graph()
sess = tf.Session()
m = model_dataspace.Model(config, name=model_uid)
_ = m() # noqa
# Create summaries
tf.summary.scalar('Train_Loss', m.vae_loss)
tf.summary.scalar('Mean_Recon_LL', m.mean_recons)
tf.summary.scalar('Mean_KL', m.mean_KL)
scalar_summaries = tf.summary.merge_all()
x_mean_, x_ = m.x_mean, m.x
if common.dataset_is_mnist_family(config['dataset']):
x_mean_ = tf.reshape(x_mean_, [-1, MNIST_SIZE, MNIST_SIZE, 1])
x_ = tf.reshape(x_, [-1, MNIST_SIZE, MNIST_SIZE, 1])
x_mean_summary = tf.summary.image(
'Reconstruction', nn.tf_batch_image(x_mean_), max_outputs=1)
x_summary = tf.summary.image('Original', nn.tf_batch_image(x_), max_outputs=1)
sample_summary = tf.summary.image(
'Sample', nn.tf_batch_image(x_mean_), max_outputs=1)
# Summary writers
train_writer = tf.summary.FileWriter(save_dir + '/vae_train', sess.graph)
eval_writer = tf.summary.FileWriter(save_dir + '/vae_eval', sess.graph)
# Initialize
sess.run(tf.global_variables_initializer())
i_start = 0
running_N_eval = 30 # pylint:disable=invalid-name
traces = {
'i': [],
'i_pred': [],
'loss': [],
'loss_eval': [],
}
best_eval_loss = np.inf
vae_lr_ = np.logspace(np.log10(FLAGS.lr), np.log10(1e-6), FLAGS.n_iters)
# Train the VAE
for i in range(i_start, FLAGS.n_iters):
start = (i * batch_size) % N_train
end = start + batch_size
batch = train_data[start:end]
labels = attr_train[start:end]
# train op
res = sess.run(
[m.train_vae, m.vae_loss, m.mean_recons, m.mean_KL, scalar_summaries], {
m.x: batch,
m.vae_lr: vae_lr_[i],
m.labels: labels,
})
tf.logging.info('Iter: %d, Loss: %d', i, res[1])
train_writer.add_summary(res[-1], i)
if i % FLAGS.n_iters_per_eval == 0:
# write training reconstructions
if batch.shape[0] == batch_size:
res = sess.run([x_summary, x_mean_summary], {
m.x: batch,
m.labels: labels,
})
train_writer.add_summary(res[0], i)
train_writer.add_summary(res[1], i)
# write sample reconstructions
prior_sample = sess.run(m.prior_sample)
res = sess.run([sample_summary], {
m.q_z_sample: prior_sample,
m.labels: labels,
})
train_writer.add_summary(res[0], i)
# write eval summaries
start = (i * batch_size) % N_eval
end = start + batch_size
batch = eval_data[start:end]
labels = attr_eval[start:end]
if batch.shape[0] == batch_size:
res_eval = sess.run([
m.vae_loss, m.mean_recons, m.mean_KL, scalar_summaries, x_summary,
x_mean_summary
], {
m.x: batch,
m.labels: labels,
})
traces['loss_eval'].append(res_eval[0])
eval_writer.add_summary(res_eval[-3], i)
eval_writer.add_summary(res_eval[-2], i)
eval_writer.add_summary(res_eval[-1], i)
if i % FLAGS.n_iters_per_save == 0:
smoothed_eval_loss = np.mean(traces['loss_eval'][-running_N_eval:])
if smoothed_eval_loss < best_eval_loss:
# Save the best model
best_eval_loss = smoothed_eval_loss
save_name = os.path.join(best_dir, 'vae_best_%s.ckpt' % model_uid)
tf.logging.info('SAVING BEST! %s Iter: %d', save_name, i)
m.vae_saver.save(sess, save_name)
with tf.gfile.Open(os.path.join(best_dir, 'best_ckpt_iters.txt'),
'w') as f:
f.write('%d' % i)
def main(unused_argv):
del unused_argv
# Load Config
config_name = FLAGS.config
config_module = importlib.import_module(configs_module_prefix +
'.%s' % config_name)
config = config_module.config
model_uid = common.get_model_uid(config_name, FLAGS.exp_uid)
batch_size = config['batch_size']
# Load dataset
dataset = common.load_dataset(config)
save_path = dataset.save_path
train_data = dataset.train_data
attr_train = dataset.attr_train
eval_data = dataset.eval_data
attr_eval = dataset.attr_eval
# Make the directory
save_dir = os.path.join(save_path, model_uid)
best_dir = os.path.join(save_dir, 'best')
tf.gfile.MakeDirs(save_dir)
tf.gfile.MakeDirs(best_dir)
tf.logging.info('Save Dir: %s', save_dir)
np.random.seed(10003)
N_train = train_data.shape[0]
N_eval = eval_data.shape[0]
# Load Model
tf.reset_default_graph()
sess = tf.Session()
m = model_dataspace.Model(config, name=model_uid)
_ = m() # noqa
# Create summaries
y_true = m.labels
y_pred = tf.cast(tf.greater(m.pred_classifier, 0.5), tf.int32)
accuracy = tf.reduce_mean(tf.cast(tf.equal(y_true, y_pred), tf.float32))
tf.summary.scalar('Loss', m.classifier_loss)
tf.summary.scalar('Accuracy', accuracy)
scalar_summaries = tf.summary.merge_all()
# Summary writers
train_writer = tf.summary.FileWriter(save_dir + '/train', sess.graph)
eval_writer = tf.summary.FileWriter(save_dir + '/eval', sess.graph)
# Initialize
sess.run(tf.global_variables_initializer())
i_start = 0
running_N_eval = 30
traces = {
'i': [],
'i_pred': [],
'loss': [],
'loss_eval': [],
}
best_eval_loss = np.inf
classifier_lr_ = np.logspace(np.log10(FLAGS.lr), np.log10(1e-6),
FLAGS.n_iters)
# Train the Classifier
for i in range(i_start, FLAGS.n_iters):
start = (i * batch_size) % N_train
end = start + batch_size
batch = train_data[start:end]
labels = attr_train[start:end]
# train op
res = sess.run(
[m.train_classifier, m.classifier_loss, scalar_summaries], {
m.x: batch,
m.labels: labels,
m.classifier_lr: classifier_lr_[i]
})
tf.logging.info('Iter: %d, Loss: %.2e', i, res[1])
train_writer.add_summary(res[-1], i)
if i % 10 == 0:
# write training reconstructions
if batch.shape[0] == batch_size:
# write eval summaries
start = (i * batch_size) % N_eval
end = start + batch_size
batch = eval_data[start:end]
labels = attr_eval[start:end]
if batch.shape[0] == batch_size:
res_eval = sess.run([m.classifier_loss, scalar_summaries],
{
m.x: batch,
m.labels: labels,
})
traces['loss_eval'].append(res_eval[0])
eval_writer.add_summary(res_eval[-1], i)
if i % FLAGS.n_iters_per_save == 0:
smoothed_eval_loss = np.mean(traces['loss_eval'][-running_N_eval:])
if smoothed_eval_loss < best_eval_loss:
# Save the best model
best_eval_loss = smoothed_eval_loss
save_name = os.path.join(best_dir,
'classifier_best_%s.ckpt' % model_uid)
tf.logging.info('SAVING BEST! %s Iter: %d', save_name, i)
m.classifier_saver.save(sess, save_name)
with tf.gfile.Open(
os.path.join(best_dir, 'best_ckpt_iters.txt'),
'w') as f:
f.write('%d' % i)
def main(unused_argv):
del unused_argv
# Load Config
config_name = FLAGS.config
config_module = importlib.import_module('configs.%s' % config_name)
config = config_module.config
model_uid = common.get_model_uid(config_name, FLAGS.exp_uid)
n_latent = config['n_latent']
# Load dataset
dataset = common.load_dataset(config)
basepath = dataset.basepath
save_path = dataset.save_path
train_data = dataset.train_data
# Make the directory
save_dir = os.path.join(save_path, model_uid)
best_dir = os.path.join(save_dir, 'best')
tf.gfile.MakeDirs(save_dir)
tf.gfile.MakeDirs(best_dir)
tf.logging.info('Save Dir: %s', save_dir)
# Set random seed
np.random.seed(FLAGS.random_seed)
tf.set_random_seed(FLAGS.random_seed)
# Load Model
tf.reset_default_graph()
sess = tf.Session()
with tf.device(tf.train.replica_device_setter(ps_tasks=0)):
m = model_dataspace.Model(config, name=model_uid)
_ = m() # noqa
# Initialize
sess.run(tf.global_variables_initializer())
# Load
m.vae_saver.restore(
sess, os.path.join(best_dir, 'vae_best_%s.ckpt' % model_uid))
# Sample from prior
sample_count = 64
image_path = os.path.join(basepath, 'sample', model_uid)
tf.gfile.MakeDirs(image_path)
# from prior
z_p = np.random.randn(sample_count, m.n_latent)
x_p = sess.run(m.x_mean, {m.z: z_p})
x_p = common.post_proc(x_p, config)
common.save_image(common.batch_image(x_p),
os.path.join(image_path, 'sample_prior.png'))
# Sample from priro, as Grid
boundary = 2.0
number_grid = 50
blob = common.make_grid(boundary=boundary,
number_grid=number_grid,
dim_latent=n_latent)
z_grid, dim_grid = blob.z_grid, blob.dim_grid
x_grid = sess.run(m.x_mean, {m.z: z_grid})
x_grid = common.post_proc(x_grid, config)
batch_image_grid = common.make_batch_image_grid(dim_grid, number_grid)
common.save_image(batch_image_grid(x_grid),
os.path.join(image_path, 'sample_grid.png'))
# Reconstruction
sample_count = 64
x_real = train_data[:sample_count]
mu, sigma = sess.run([m.mu, m.sigma], {m.x: x_real})
x_rec = sess.run(m.x_mean, {m.mu: mu, m.sigma: sigma})
x_rec = common.post_proc(x_rec, config)
x_real = common.post_proc(x_real, config)
common.save_image(common.batch_image(x_real),
os.path.join(image_path, 'image_real.png'))
common.save_image(common.batch_image(x_rec),
os.path.join(image_path, 'image_rec.png'))
def main(unused_argv):
del unused_argv
# Load Config
config_name = FLAGS.config
config_module = importlib.import_module('configs.%s' % config_name)
config = config_module.config
model_uid = common.get_model_uid(config_name, FLAGS.exp_uid)
batch_size = config['batch_size']
# Load dataset
dataset = common.load_dataset(config)
basepath = dataset.basepath
save_path = dataset.save_path
train_data = dataset.train_data
eval_data = dataset.eval_data
# Make the directory
save_dir = os.path.join(save_path, model_uid)
best_dir = os.path.join(save_dir, 'best')
tf.gfile.MakeDirs(save_dir)
tf.gfile.MakeDirs(best_dir)
tf.logging.info('Save Dir: %s', save_dir)
# Load Model
tf.reset_default_graph()
sess = tf.Session()
m = model_dataspace.Model(config, name=model_uid)
_ = m() # noqa
# Initialize
sess.run(tf.global_variables_initializer())
# Load
m.vae_saver.restore(sess,
os.path.join(best_dir, 'vae_best_%s.ckpt' % model_uid))
# Encode
def encode(data):
"""Encode the data in dataspace to latent spaceself.
This script runs the encoding in batched mode to limit GPU memory usage.
Args:
data: A numpy array of data to be encoded.
Returns:
A object with instances `mu` and `sigma`, the parameters of encoded
distributions in the latent space.
"""
mu_list, sigma_list = [], []
for i in range(0, len(data), batch_size):
start, end = i, min(i + batch_size, len(data))
batch = data[start:end]
mu, sigma = sess.run([m.mu, m.sigma], {m.x: batch})
mu_list.append(mu)
sigma_list.append(sigma)
mu = np.concatenate(mu_list)
sigma = np.concatenate(sigma_list)
return common.ObjectBlob(mu=mu, sigma=sigma)
encoded_train_data = encode(train_data)
tf.logging.info(
'encode train_data: mu.shape = %s sigma.shape = %s',
encoded_train_data.mu.shape,
encoded_train_data.sigma.shape,
)
encoded_eval_data = encode(eval_data)
tf.logging.info(
'encode eval_data: mu.shape = %s sigma.shape = %s',
encoded_eval_data.mu.shape,
encoded_eval_data.sigma.shape,
)
# Save encoded as npz file
encoded_save_path = os.path.join(basepath, 'encoded', model_uid)
tf.gfile.MakeDirs(encoded_save_path)
tf.logging.info('encoded train_data saved to %s',
os.path.join(encoded_save_path, 'encoded_train_data.npz'))
np.savez(
os.path.join(encoded_save_path, 'encoded_train_data.npz'),
mu=encoded_train_data.mu,
sigma=encoded_train_data.sigma,
)
tf.logging.info('encoded eval_data saved to %s',
os.path.join(encoded_save_path, 'encoded_eval_data.npz'))
np.savez(
os.path.join(encoded_save_path, 'encoded_eval_data.npz'),
mu=encoded_eval_data.mu,
sigma=encoded_eval_data.sigma,
)