def while_body(g_pools, g_logits, i):
with tf.control_dependencies([g_pools, g_logits]):
test_zs = utils.make_z_normal(1, local_batch_size, FLAGS.z_dim)
# Uniform distribution
# TODO(goodfellow) Use true distribution of ImageNet classses
gen_class_logits = tf.zeros((local_batch_size, num_classes))
gen_class_ints = tf.multinomial(gen_class_logits, 1)
gen_sparse_class = tf.squeeze(gen_class_ints)
with tf.variable_scope('model'):
generator = generator_fn(
test_zs[0],
gen_sparse_class,
FLAGS.gf_dim,
FLAGS.num_classes)
pools, logits = utils.run_custom_inception(
generator,
output_tensor=['pool_3:0', 'logits:0'],
graph_def=graph_def)
g_pools = tf.concat([g_pools, pools], 0)
g_logits = tf.concat([g_logits, logits], 0)
return (g_pools, g_logits, tf.add(i, 1))
def while_body(g_pools, g_logits, i):
with tf.control_dependencies([g_pools, g_logits]):
test_zs = utils.make_z_normal(1, local_batch_size, FLAGS.z_dim)
# Uniform distribution
gen_class_logits = tf.zeros((local_batch_size, num_classes))
gen_class_ints = tf.multinomial(gen_class_logits, 1)
gen_sparse_class = tf.squeeze(gen_class_ints)
generator = generator_fn(test_zs[0],
gen_sparse_class,
FLAGS.gf_dim,
FLAGS.num_classes,
is_training=False)
pools, logits = utils.run_custom_inception(
generator,
output_tensor=['pool_3:0', 'logits:0'],
graph_def=graph_def)
g_pools = tf.concat([g_pools, pools], 0)
g_logits = tf.concat([g_logits, logits], 0)
return (g_pools, g_logits, tf.add(i, 1))
def main(_, is_test=False):
print('d_learning_rate', FLAGS.discriminator_learning_rate)
print('g_learning_rate', FLAGS.generator_learning_rate)
print('data_dir', FLAGS.data_dir)
print(FLAGS.loss_type, FLAGS.batch_size, FLAGS.beta1)
print('gf_df_dim', FLAGS.gf_dim, FLAGS.df_dim)
print('Starting the program..')
gfile.MakeDirs(FLAGS.checkpoint_dir)
model_dir = '%s_%s' % ('celebA', FLAGS.batch_size)
logdir = os.path.join(FLAGS.checkpoint_dir, model_dir)
gfile.MakeDirs(logdir)
graph = tf.Graph()
with graph.as_default():
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Instantiate global_step.
global_step = tf.train.create_global_step()
# Create model with FLAGS, global_step, and devices.
devices = [
'/gpu:{}'.format(tower) for tower in range(FLAGS.num_towers)
]
# Create noise tensors
zs = utils.make_z_normal(FLAGS.num_towers, FLAGS.batch_size,
FLAGS.z_dim)
print('save_summaries_steps', FLAGS.save_summaries_steps)
dcgan = model.SNGAN(zs=zs,
config=FLAGS,
global_step=global_step,
devices=devices)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Create sync_hooks when needed.
if FLAGS.sync_replicas and FLAGS.num_workers > 1:
print('condition 1')
sync_hooks = [
dcgan.d_opt.make_session_run_hook(FLAGS.task == 0),
dcgan.g_opt.make_session_run_hook(FLAGS.task == 0)
]
else:
print('condition 2')
sync_hooks = []
train_ops = tfgan.GANTrainOps(
generator_train_op=dcgan.g_optim,
discriminator_train_op=dcgan.d_optim,
global_step_inc_op=dcgan.increment_global_step)
# We set allow_soft_placement to be True because Saver for the DCGAN model
# gets misplaced on the GPU.
session_config = tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False)
if is_test:
return graph
print("G step: ", FLAGS.g_step)
print("D_step: ", FLAGS.d_step)
train_steps = tfgan.GANTrainSteps(FLAGS.g_step, FLAGS.d_step)
tfgan.gan_train(
train_ops,
get_hooks_fn=tfgan.get_sequential_train_hooks(
train_steps=train_steps),
hooks=([tf.train.StopAtStepHook(num_steps=2000000)] + sync_hooks),
logdir=logdir,
# master=FLAGS.master,
# scaffold=scaffold, # load from google checkpoint
is_chief=(FLAGS.task == 0),
save_summaries_steps=FLAGS.save_summaries_steps,
save_checkpoint_secs=FLAGS.save_checkpoint_secs,
config=session_config)
def main(_):
model_dir = '%s_%s' % ('imagenet', FLAGS.batch_size)
checkpoint_dir = os.path.join(FLAGS.checkpoint_dir, model_dir)
log_dir = os.path.join(FLAGS.eval_dir, model_dir)
graph_def = None # pylint: disable=protected-access
# Batch size to feed batches of images through Inception and the generator
# to extract feature vectors to later stack together and compute metrics.
local_batch_size = FLAGS.dcgan_generator_batch_size
if FLAGS.generator_type == 'baseline':
generator_fn = generator_module.generator
elif FLAGS.generator_type == 'test':
generator_fn = generator_module.generator_test
else:
raise NotImplementedError
if FLAGS.num_towers != 1 or FLAGS.num_workers != 1:
raise NotImplementedError(
'The eval job does not currently support using multiple GPUs')
# Get activations from real images.
with tf.device('/device:CPU:1'):
real_pools, real_images = utils.get_real_activations(
FLAGS.data_dir,
local_batch_size,
FLAGS.eval_sample_size // local_batch_size,
label_offset=-1,
shuffle_buffer_size=FLAGS.shuffle_buffer_size)
# Uniform distribution
# TODO(goodfellow) Use true distribution of ImageNet classses
num_classes = FLAGS.num_classes
gen_class_logits = tf.zeros((local_batch_size, num_classes))
gen_class_ints = tf.multinomial(gen_class_logits, 1)
gen_sparse_class = tf.squeeze(gen_class_ints)
with tf.variable_scope('model'):
# Generate the first batch of generated images and extract activations;
# this bootstraps the while_loop with a pools and logits tensor.
test_zs = utils.make_z_normal(1, local_batch_size, FLAGS.z_dim)
generator = generator_fn(
test_zs[0],
gen_sparse_class,
FLAGS.gf_dim,
FLAGS.num_classes)
pools, logits = utils.run_custom_inception(
generator, output_tensor=['pool_3:0', 'logits:0'], graph_def=graph_def)
# Set up while_loop to compute activations of generated images from generator.
def while_cond(g_pools, g_logits, i): # pylint: disable=unused-argument
return tf.less(i, FLAGS.eval_sample_size // local_batch_size)
# We use a while loop because we want to generate a batch of images
# and then feed that batch through Inception to retrieve the activations.
# Otherwise, if we generate all the samples first and then compute all the
# activations, we will run out of memory.
def while_body(g_pools, g_logits, i):
with tf.control_dependencies([g_pools, g_logits]):
test_zs = utils.make_z_normal(1, local_batch_size, FLAGS.z_dim)
# Uniform distribution
# TODO(goodfellow) Use true distribution of ImageNet classses
gen_class_logits = tf.zeros((local_batch_size, num_classes))
gen_class_ints = tf.multinomial(gen_class_logits, 1)
gen_sparse_class = tf.squeeze(gen_class_ints)
with tf.variable_scope('model'):
generator = generator_fn(
test_zs[0],
gen_sparse_class,
FLAGS.gf_dim,
FLAGS.num_classes)
pools, logits = utils.run_custom_inception(
generator,
output_tensor=['pool_3:0', 'logits:0'],
graph_def=graph_def)
g_pools = tf.concat([g_pools, pools], 0)
g_logits = tf.concat([g_logits, logits], 0)
return (g_pools, g_logits, tf.add(i, 1))
# Get the activations
i = tf.constant(1)
new_generator_pools_list, new_generator_logits_list, _ = tf.while_loop(
while_cond,
while_body, [pools, logits, i],
shape_invariants=[
tf.TensorShape([None, 2048]),
tf.TensorShape([None, 1008]),
i.get_shape()
],
parallel_iterations=1,
back_prop=False,
swap_memory=True,
name='GeneratedActivations')
new_generator_pools_list.set_shape([FLAGS.eval_sample_size, 2048])
new_generator_logits_list.set_shape([FLAGS.eval_sample_size, 1008])
# TODO(sbhupatiraju) Why is FID negative?
# Get a small batch of samples from generator to dispaly in TensorBoard
vis_batch_size = 16
eval_vis_zs = utils.make_z_normal(
1, vis_batch_size, FLAGS.z_dim)
# Uniform distribution
# TODO(goodfellow) Use true distribution of ImageNet classses
gen_class_logits_vis = tf.zeros((vis_batch_size, num_classes))
gen_class_ints_vis = tf.multinomial(gen_class_logits_vis, 1)
gen_sparse_class_vis = tf.squeeze(gen_class_ints_vis)
with tf.variable_scope('model'):
eval_vis_images = generator_fn(
eval_vis_zs[0],
gen_sparse_class_vis,
FLAGS.gf_dim,
FLAGS.num_classes
)
eval_vis_images = tf.cast((eval_vis_images + 1.) * 127.5, tf.uint8)
with tf.variable_scope('eval_vis'):
tf.summary.image('generated_images', eval_vis_images)
tf.summary.image('real_images', real_images)
tf.summary.image('real_images_grid',
tfgan.eval.image_grid(
real_images[:16],
grid_shape=utils.squarest_grid_size(16),
image_shape=(128, 128)))
tf.summary.image('generated_images_grid',
tfgan.eval.image_grid(
eval_vis_images[:16],
grid_shape=utils.squarest_grid_size(16),
image_shape=(128, 128)))
# Use the activations from the real images and generated images to compute
# Inception score and FID.
generated_logits = tf.concat(new_generator_logits_list, 0)
generated_pools = tf.concat(new_generator_pools_list, 0)
# Compute Frechet Inception Distance and Inception score
incscore = tfgan.eval.classifier_score_from_logits(generated_logits)
fid = tfgan.eval.frechet_classifier_distance_from_activations(
real_pools, generated_pools)
with tf.variable_scope('eval'):
tf.summary.scalar('fid', fid)
tf.summary.scalar('incscore', incscore)
session_config = tf.ConfigProto(
allow_soft_placement=True, log_device_placement=False)
tf.contrib.training.evaluate_repeatedly(
checkpoint_dir=checkpoint_dir,
hooks=[
tf.contrib.training.SummaryAtEndHook(log_dir),
tf.contrib.training.StopAfterNEvalsHook(1)
],
config=session_config)