def precompute_gram_matrices(image, final_endpoint='fc8'):
"""Pre-computes the Gram matrices on a given image.
Args:
image: 4-D tensor. Input (batch of) image(s).
final_endpoint: str, name of the final layer to compute Gram matrices for.
Defaults to 'fc8'.
Returns:
dict mapping layer names to their corresponding Gram matrices.
"""
with tf.Session() as session:
end_points = vgg.vgg_16(image, final_endpoint=final_endpoint)
tf.train.Saver(slim.get_variables('vgg_16')).restore(
session, vgg.checkpoint_file())
return dict((key, gram_matrix(value).eval())
for key, value in end_points.items())
def precompute_gram_matrices(image, final_endpoint='fc8'):
"""Pre-computes the Gram matrices on a given image.
Args:
image: 4-D tensor. Input (batch of) image(s).
final_endpoint: str, name of the final layer to compute Gram matrices for.
Defaults to 'fc8'.
Returns:
dict mapping layer names to their corresponding Gram matrices.
"""
with tf.Session() as session:
end_points = vgg.vgg_16(image, final_endpoint=final_endpoint)
tf.train.Saver(slim.get_variables('vgg_16')).restore(
session, vgg.checkpoint_file())
return dict([(key, _gram_matrix(value).eval())
for key, value in end_points.iteritems()])
def main(unused_argv=None):
with tf.Graph().as_default():
# Force all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks,
worker_device=device)):
inputs, _ = image_utils.imagenet_inputs(FLAGS.batch_size,
FLAGS.image_size)
# Load style images and select one at random (for each graph execution, a
# new random selection occurs)
style_images, style_labels, \
style_gram_matrices = image_utils.style_image_inputs(
os.path.expanduser(FLAGS.style_dataset_file),
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
square_crop=True,
shuffle=True)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and weight flags
num_styles = FLAGS.num_styles
if FLAGS.style_coefficients is None:
style_coefficients = [1.0 for _ in range(num_styles)]
else:
style_coefficients = ast.literal_eval(FLAGS.style_coefficients)
if len(style_coefficients) != num_styles:
raise ValueError(
'number of style coefficients differs from number of styles'
)
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Rescale style weights dynamically based on the current style image
style_coefficient = tf.gather(tf.constant(style_coefficients),
style_labels)
style_weights = dict((key, style_coefficient * style_weights[key])
for key in style_weights)
# Define the model
stylized_inputs = model.transform(inputs,
alpha=FLAGS.alpha,
normalizer_params={
'labels': style_labels,
'num_categories': num_styles,
'center': True,
'scale': True
})
# Compute losses.
total_loss, loss_dict = learning.total_loss(
inputs, stylized_inputs, style_gram_matrices, content_weights,
style_weights)
for key in loss_dict:
tf.summary.scalar(key, loss_dict[key])
# Adding Image summaries to the tensorboard.
tf.summary.image('image/0_inputs', inputs, 3)
tf.summary.image('image/1_styles', style_images, 3)
tf.summary.image('image/2_styled_inputs', stylized_inputs, 3)
# Set up training
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
summarize_gradients=False)
# Function to restore VGG16 parameters.
init_fn_vgg = slim.assign_from_checkpoint_fn(
vgg.checkpoint_file(), slim.get_variables('vgg_16'))
# Run training
slim.learning.train(train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_fn_vgg,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
def init_fn_vgg(session):
saver_vgg.restore(session, vgg.checkpoint_file())
def main(unused_argv=None):
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
# Forces all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks,
worker_device=device)):
# Loads content images.
content_inputs_, _ = image_utils.imagenet_inputs(
FLAGS.batch_size, FLAGS.image_size)
# Loads style images.
[style_inputs_, _, _] = image_utils.arbitrary_style_image_inputs(
FLAGS.style_dataset_file,
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
shuffle=True,
center_crop=FLAGS.center_crop,
augment_style_images=FLAGS.augment_style_images,
random_style_image_size=FLAGS.random_style_image_size)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and content weight flags.
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Define the model
stylized_images, total_loss, loss_dict, \
_ = build_mobilenet_model.build_mobilenet_model(
content_inputs_,
style_inputs_,
mobilenet_trainable=False,
style_params_trainable=True,
transformer_trainable=True,
mobilenet_end_point='layer_19',
transformer_alpha=FLAGS.alpha,
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=FLAGS.total_variation_weight,
)
# Adding scalar summaries to the tensorboard.
for key in loss_dict:
tf.summary.scalar(key, loss_dict[key])
# Adding Image summaries to the tensorboard.
tf.summary.image('image/0_content_inputs', content_inputs_, 3)
tf.summary.image('image/1_style_inputs_aug', style_inputs_, 3)
tf.summary.image('image/2_stylized_images', stylized_images, 3)
# Set up training
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
summarize_gradients=False)
# Function to restore VGG16 parameters.
init_fn_vgg = slim.assign_from_checkpoint_fn(
vgg.checkpoint_file(), slim.get_variables('vgg_16'))
# Function to restore Mobilenet V2 parameters.
mobilenet_variables_dict = {
var.op.name: var
for var in slim.get_model_variables('MobilenetV2')
}
init_fn_mobilenet = slim.assign_from_checkpoint_fn(
FLAGS.mobilenet_checkpoint, mobilenet_variables_dict)
# Function to restore VGG16 and Mobilenet V2 parameters.
def init_sub_networks(session):
init_fn_vgg(session)
init_fn_mobilenet(session)
# Run training
slim.learning.train(train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_sub_networks,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
def init_fn_vgg(session): saver_vgg.restore(session, vgg.checkpoint_file())
def main(unused_argv=None):
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
# Forces all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks,
worker_device=device)):
# Loads content images.
# content_inputs_, _ = image_utils.imagenet_inputs(FLAGS.batch_size,
# FLAGS.image_size)
#changed
content_inputs_, _, temp = load_xy_pairs(
fake_directory=FLAGS.content_dir,
real_directory=FLAGS.content_dir,
batch_size=FLAGS.batch_size,
prob_of_real=1)
# Loads style images.
# [style_inputs_, _,
# style_inputs_orig_] = image_utils.arbitrary_style_image_inputs(
# FLAGS.style_dataset_file,
# batch_size=FLAGS.batch_size,
# image_size=FLAGS.image_size,
# shuffle=True,
# center_crop=FLAGS.center_crop,
# augment_style_images=FLAGS.augment_style_images,
# random_style_image_size=FLAGS.random_style_image_size)
style_inputs_, _, temp2 = load_xy_pairs(
fake_directory=FLAGS.style_dir,
real_directory=FLAGS.style_dir,
batch_size=FLAGS.batch_size,
prob_of_real=1)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and content weight flags.
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Define the model
stylized_images, total_loss, loss_dict, _ = build_model.build_model(
content_inputs_,
style_inputs_,
trainable=True,
is_training=True,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=FLAGS.total_variation_weight)
if tf.gfile.IsDirectory(FLAGS.checkpoint):
checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoint)
else:
checkpoint = FLAGS.checkpoint
tf.logging.info(
'loading latest checkpoint file: {}'.format(checkpoint))
init_fn = slim.assign_from_checkpoint_fn(
checkpoint, slim.get_variables_to_restore())
# sess.run([tf.local_variables_initializer()])
# init_fn(sess)
# Adding scalar summaries to the tensorboard.
for key, value in loss_dict.iteritems():
tf.summary.scalar(key, value)
# Adding Image summaries to the tensorboard.
tf.summary.image('image/0_content_inputs', content_inputs_, 3)
# tf.summary.image('image/1_style_inputs_orig', style_inputs_orig_, 3)
tf.summary.image('image/2_style_inputs_aug', style_inputs_, 3)
tf.summary.image('image/3_stylized_images', stylized_images, 3)
# Set up training
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
summarize_gradients=False)
# Function to restore VGG16 parameters.
init_fn_vgg = slim.assign_from_checkpoint_fn(
vgg.checkpoint_file(), slim.get_variables('vgg_16'))
# Function to restore Inception_v3 parameters.
inception_variables_dict = {
var.op.name: var
for var in slim.get_model_variables('InceptionV3')
}
init_fn_inception = slim.assign_from_checkpoint_fn(
FLAGS.inception_v3_checkpoint, inception_variables_dict)
# Function to restore VGG16 and Inception_v3 parameters.
def init_sub_networks(session):
# init_fn_vgg(session)
# init_fn_inception(session)
# session.run(tf.local_variables_initializer())
init_fn(session)
# config = tf.ConfigProto()
# config.gpu_options_allow_growth = True
config = tf.ConfigProto()
# config.gpu_options.allocator_type = 'BFC'
# config.gpu_options.per_process_gpu_memory_fraction = 0.98
# Run training
slim.learning.train(session_config=config,
train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_sub_networks,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
def main(unused_argv=None):
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
# Forces all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks, worker_device=device)):
# Loads content images.
content_inputs_, _ = image_utils.imagenet_inputs(FLAGS.batch_size,
FLAGS.image_size)
# Loads style images.
[style_inputs_, _,
style_inputs_orig_] = image_utils.arbitrary_style_image_inputs(
FLAGS.style_dataset_file,
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
shuffle=True,
center_crop=FLAGS.center_crop,
augment_style_images=FLAGS.augment_style_images,
random_style_image_size=FLAGS.random_style_image_size)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and content weight flags.
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Define the model
stylized_images, total_loss, loss_dict, _ = build_model.build_model(
content_inputs_,
style_inputs_,
trainable=True,
is_training=True,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=FLAGS.total_variation_weight)
# Adding scalar summaries to the tensorboard.
for key, value in loss_dict.iteritems():
tf.summary.scalar(key, value)
# Adding Image summaries to the tensorboard.
tf.summary.image('image/0_content_inputs', content_inputs_, 3)
tf.summary.image('image/1_style_inputs_orig', style_inputs_orig_, 3)
tf.summary.image('image/2_style_inputs_aug', style_inputs_, 3)
tf.summary.image('image/3_stylized_images', stylized_images, 3)
# Set up training
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
summarize_gradients=False)
# Function to restore VGG16 parameters.
init_fn_vgg = slim.assign_from_checkpoint_fn(vgg.checkpoint_file(),
slim.get_variables('vgg_16'))
# Function to restore Inception_v3 parameters.
inception_variables_dict = {
var.op.name: var
for var in slim.get_model_variables('InceptionV3')
}
init_fn_inception = slim.assign_from_checkpoint_fn(
FLAGS.inception_v3_checkpoint, inception_variables_dict)
# Function to restore VGG16 and Inception_v3 parameters.
def init_sub_networks(session):
init_fn_vgg(session)
init_fn_inception(session)
# Run training
slim.learning.train(
train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_sub_networks,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
def main(unused_argv=None):
with tf.Graph().as_default():
# Force all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks,
worker_device=device)):
inputs, _ = image_utils.imagenet_inputs(FLAGS.batch_size,
FLAGS.image_size)
# Load style images and select one at random (for each graph execution, a
# new random selection occurs)
_, style_labels, \
style_gram_matrices = image_utils.style_image_inputs(
os.path.expanduser(FLAGS.style_dataset_file),
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
square_crop=True,
shuffle=True)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and weight flags
num_styles = FLAGS.num_styles
if FLAGS.style_coefficients is None:
style_coefficients = [1.0 for _ in range(num_styles)]
else:
style_coefficients = ast.literal_eval(FLAGS.style_coefficients)
if len(style_coefficients) != num_styles:
raise ValueError(
'number of style coefficients differs from number of styles'
)
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Rescale style weights dynamically based on the current style image
style_coefficient = tf.gather(tf.constant(style_coefficients),
style_labels)
style_weights = dict((key, style_coefficient * value)
for key, value in style_weights.items())
# Define the model
stylized_inputs = model.transform(inputs,
alpha=FLAGS.alpha,
normalizer_params={
'labels': style_labels,
'num_categories': num_styles,
'center': True,
'scale': True
})
# Compute losses.
total_loss, loss_dict = learning.total_loss(
inputs, stylized_inputs, style_gram_matrices, content_weights,
style_weights)
for key, value in loss_dict.items():
tf.summary.scalar(key, value)
instance_norm_vars = [
var for var in slim.get_variables('transformer')
if 'InstanceNorm' in var.name
]
other_vars = [
var for var in slim.get_variables('transformer')
if 'InstanceNorm' not in var.name
]
# Function to restore VGG16 parameters.
init_fn_vgg = slim.assign_from_checkpoint_fn(
vgg.checkpoint_file(), slim.get_variables('vgg_16'))
checkpoint = os.path.expanduser(FLAGS.checkpoint)
if tf.gfile.IsDirectory(checkpoint):
checkpoint = tf.train.latest_checkpoint(checkpoint)
tf.logging.info(
'loading latest checkpoint file: {}'.format(checkpoint))
# Function to restore N-styles parameters.
vars = slim.get_variables(
'transformer') if FLAGS.restore_all_weights else other_vars
init_fn_n_styles = slim.assign_from_checkpoint_fn(checkpoint, vars)
def init_fn(session):
init_fn_vgg(session)
init_fn_n_styles(session)
# Set up training.
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
variables_to_train=instance_norm_vars,
summarize_gradients=False)
savertransformer = tf.train.Saver(
variables.get_variables("transformer"),
save_relative_paths=True)
# Run training.
slim.learning.train(train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
log_every_n_steps=FLAGS.log_every_n_steps,
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_fn,
saver=savertransformer,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
