def generate(network_pkl, out_dir):
if os.path.exists(out_dir):
raise ValueError('{} already exists'.format(out_dir))
misc.init_output_logging()
np.random.seed(config.random_seed)
tfutil.init_tf(config.tf_config)
with tf.device('/gpu:0'):
G, D, Gs = misc.load_pkl(network_pkl)
training_set = dataset.load_dataset(data_dir=config.data_dir,
verbose=True,
**config.dataset)
# grid_size, grid_reals, grid_labels, grid_latents = train.setup_snapshot_image_grid(G, training_set, **config.grid)
number_of_images = 1000
grid_labels = np.zeros([number_of_images, training_set.label_size],
dtype=training_set.label_dtype)
grid_latents = misc.random_latents(number_of_images, G)
total_kimg = config.train.total_kimg
sched = train.TrainingSchedule(total_kimg * 1000, training_set,
**config.sched)
grid_fakes = Gs.run(grid_latents,
grid_labels,
minibatch_size=sched.minibatch // config.num_gpus)
os.makedirs(out_dir)
# print(np.min(grid_fakes), np.mean(grid_fakes), np.max(grid_fakes))
# misc.save_image_grid(grid_fakes, 'fakes.png', drange=[-1,1], grid_size=grid_size)
for i, img in enumerate(grid_fakes):
img = img.transpose((1, 2, 0))
img = np.clip(img, -1, 1)
img = (1 + img) / 2
img = skimage.img_as_ubyte(img)
imageio.imwrite(os.path.join(out_dir, '{}.png'.format(i)),
img[..., :3])
if img.shape[-1] > 3:
np.save(os.path.join(out_dir, '{}.npy'.format(i)), img)
(G, D, Gs, E),
os.path.join(
result_subdir,
'network-snapshot-%06d.pkl' % (cur_nimg // 1000)))
# Record start time of the next tick.
tick_start_time = time.time()
# Write final results.
misc.save_pkl((G, D, Gs, E),
os.path.join(result_subdir, 'network-final.pkl'))
summary_log.close()
open(os.path.join(result_subdir, '_training-done.txt'), 'wt').close()
#----------------------------------------------------------------------------
# Main entry point.
# Calls the function indicated in config.py.
if __name__ == "__main__":
misc.init_output_logging()
np.random.seed(config.random_seed)
print('Initializing TensorFlow...')
os.environ.update(config.env)
tfutil.init_tf(config.tf_config)
print('Running %s()...' % config.train['func'])
tfutil.call_func_by_name(**config.train)
print('Exiting...')
#----------------------------------------------------------------------------
# Save snapshots.
if cur_tick % image_snapshot_ticks == 0 or done:
grid_fakes = Gs.run(grid_latents, grid_labels, minibatch_size=sched.minibatch//config.num_gpus)
misc.save_image_grid(grid_fakes, os.path.join(result_subdir, 'fakes%06d.png' % (cur_nimg // 1000)), drange=drange_net, grid_size=grid_size)
if cur_tick % network_snapshot_ticks == 0 or done:
misc.save_pkl((G, D, Gs), os.path.join(result_subdir, 'network-snapshot-%06d.pkl' % (cur_nimg // 1000)))
# Record start time of the next tick.
tick_start_time = time.time()
# Write final results.
misc.save_pkl((G, D, Gs), os.path.join(result_subdir, 'network-final.pkl'))
summary_log.close()
open(os.path.join(result_subdir, '_training-done.txt'), 'wt').close()
#----------------------------------------------------------------------------
# Main entry point.
# Calls the function indicated in config.py.
if __name__ == "__main__":
misc.init_output_logging()
np.random.seed(config.random_seed)
print('Initializing TensorFlow...')
os.environ.update(config.env)
tfutil.init_tf(config.tf_config)
print('Running %s()...' % config.train['func'])
tfutil.call_func_by_name(**config.train)
print('Exiting...')
#----------------------------------------------------------------------------
