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 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)):
# Load 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, \
true_loss, \
_, \
bottleneck_feat = build_mobilenet_model.build_mobilenet_model(
content_inputs_,
style_inputs_,
mobilenet_trainable=True,
style_params_trainable=False,
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=FLAGS.total_variation_weight,
)
_, inception_bottleneck_feat = build_model.style_prediction(
style_inputs_,
[],
[],
is_training=False,
trainable=False,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
reuse=None,
)
print('PRINTING TRAINABLE VARIABLES')
for x in tf.trainable_variables():
print(x)
mse_loss = tf.losses.mean_squared_error(inception_bottleneck_feat,
bottleneck_feat)
total_loss = mse_loss
if FLAGS.use_true_loss:
true_loss = FLAGS.true_loss_weight * true_loss
total_loss += true_loss
if FLAGS.use_true_loss:
tf.summary.scalar('mse', mse_loss)
tf.summary.scalar('true_loss', true_loss)
tf.summary.scalar('total_loss', total_loss)
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)
mobilenet_variables_to_restore = slim.get_variables_to_restore(
include=['MobilenetV2'], exclude=['global_step'])
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)
init_fn = slim.assign_from_checkpoint_fn(
FLAGS.initial_checkpoint,
slim.get_variables_to_restore(
exclude=['MobilenetV2', 'mobilenet_conv', 'global_step']))
init_pretrained_mobilenet = slim.assign_from_checkpoint_fn(
FLAGS.mobilenet_checkpoint, mobilenet_variables_to_restore)
def init_sub_networks(session):
init_fn(session)
init_pretrained_mobilenet(session)
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,
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.
# TODO(iansimon): This is ugly, but assign_from_checkpoint_fn doesn't
# exist yet.
saver_vgg = tf.train.Saver(slim.get_variables('vgg_16'))
def init_fn_vgg(session):
saver_vgg.restore(session, vgg.checkpoint_file())
# Function to restore N-styles parameters.
# TODO(iansimon): This is ugly, but assign_from_checkpoint_fn doesn't
# exist yet.
saver_n_styles = tf.train.Saver(other_vars)
def init_fn_n_styles(session):
saver_n_styles.restore(session,
os.path.expanduser(FLAGS.checkpoint))
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)
# 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,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
# Loads content images.
eval_content_inputs_, _ = image_utils.imagenet_inputs(
FLAGS.batch_size, FLAGS.image_size)
# Process style and content weight flags.
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Loads evaluation style images.
eval_style_inputs_, _, _ = image_utils.arbitrary_style_image_inputs(
FLAGS.eval_style_dataset_file,
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
center_crop=True,
shuffle=True,
augment_style_images=False,
random_style_image_size=False)
# Computes stylized noise.
stylized_noise, _, _, _ = build_model.build_model(
tf.random_uniform([
min(4, FLAGS.batch_size), FLAGS.image_size, FLAGS.image_size, 3
]),
tf.slice(eval_style_inputs_, [0, 0, 0, 0],
[min(4, FLAGS.batch_size), -1, -1, -1]),
trainable=False,
is_training=False,
reuse=None,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
adds_losses=False)
# Computes stylized images.
stylized_images, _, loss_dict, _ = build_model.build_model(
eval_content_inputs_,
eval_style_inputs_,
trainable=False,
is_training=False,
reuse=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)
# Adds Image summaries to the tensorboard.
tf.summary.image(
'image/{}/0_eval_content_inputs'.format(FLAGS.eval_name),
eval_content_inputs_, 3)
tf.summary.image(
'image/{}/1_eval_style_inputs'.format(FLAGS.eval_name),
eval_style_inputs_, 3)
tf.summary.image(
'image/{}/2_eval_stylized_images'.format(FLAGS.eval_name),
stylized_images, 3)
tf.summary.image('image/{}/3_stylized_noise'.format(FLAGS.eval_name),
stylized_noise, 3)
metrics = {}
for key, value in loss_dict.items():
metrics[key] = tf.metrics.mean(value)
names_values, names_updates = slim.metrics.aggregate_metric_map(
metrics)
for name, value in names_values.items():
slim.summaries.add_scalar_summary(value, name, print_summary=True)
eval_op = list(names_updates.values())
num_evals = FLAGS.num_evaluation_styles / FLAGS.batch_size
slim.evaluation.evaluation_loop(
master=FLAGS.master,
checkpoint_dir=FLAGS.checkpoint_dir,
logdir=FLAGS.eval_dir,
eval_op=eval_op,
num_evals=num_evals,
eval_interval_secs=FLAGS.eval_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_, _, _] = 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 main(_):
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
# Loads content images.
eval_content_inputs_, _ = image_utils.imagenet_inputs(
FLAGS.batch_size, FLAGS.image_size)
# Process style and content weight flags.
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Loads evaluation style images.
eval_style_inputs_, _, _ = image_utils.arbitrary_style_image_inputs(
FLAGS.eval_style_dataset_file,
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
center_crop=True,
shuffle=True,
augment_style_images=False,
random_style_image_size=False)
# Computes stylized noise.
stylized_noise, _, _, _ = build_model.build_model(
tf.random_uniform(
[min(4, FLAGS.batch_size), FLAGS.image_size, FLAGS.image_size, 3]),
tf.slice(eval_style_inputs_, [0, 0, 0, 0],
[min(4, FLAGS.batch_size), -1, -1, -1]),
trainable=False,
is_training=False,
reuse=None,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
adds_losses=False)
# Computes stylized images.
stylized_images, _, loss_dict, _ = build_model.build_model(
eval_content_inputs_,
eval_style_inputs_,
trainable=False,
is_training=False,
reuse=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)
# Adds Image summaries to the tensorboard.
tf.summary.image('image/{}/0_eval_content_inputs'.format(FLAGS.eval_name),
eval_content_inputs_, 3)
tf.summary.image('image/{}/1_eval_style_inputs'.format(FLAGS.eval_name),
eval_style_inputs_, 3)
tf.summary.image('image/{}/2_eval_stylized_images'.format(FLAGS.eval_name),
stylized_images, 3)
tf.summary.image('image/{}/3_stylized_noise'.format(FLAGS.eval_name),
stylized_noise, 3)
metrics = {}
for key, value in loss_dict.iteritems():
metrics[key] = tf.metrics.mean(value)
names_values, names_updates = slim.metrics.aggregate_metric_map(metrics)
for name, value in names_values.iteritems():
slim.summaries.add_scalar_summary(value, name, print_summary=True)
eval_op = names_updates.values()
num_evals = FLAGS.num_evaluation_styles / FLAGS.batch_size
slim.evaluation.evaluation_loop(
master=FLAGS.master,
checkpoint_dir=FLAGS.checkpoint_dir,
logdir=FLAGS.eval_dir,
eval_op=eval_op,
num_evals=num_evals,
eval_interval_secs=FLAGS.eval_interval_secs)
def main(_):
with tf.Graph().as_default():
# Create inputs in [0, 1], as expected by vgg_16.
inputs, _ = image_utils.imagenet_inputs(FLAGS.batch_size,
FLAGS.image_size)
evaluation_images = image_utils.load_evaluation_images(
FLAGS.image_size)
# Process style and weight flags
if FLAGS.style_coefficients is None:
style_coefficients = [1.0 for _ in range(FLAGS.num_styles)]
else:
style_coefficients = ast.literal_eval(FLAGS.style_coefficients)
if len(style_coefficients) != FLAGS.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)
# Load style images.
style_images, labels, style_gram_matrices = image_utils.style_image_inputs(
os.path.expanduser(FLAGS.style_dataset_file),
batch_size=FLAGS.num_styles,
image_size=FLAGS.image_size,
square_crop=True,
shuffle=False)
labels = tf.unstack(labels)
def _create_normalizer_params(style_label):
"""Creates normalizer parameters from a style label."""
return {
'labels': tf.expand_dims(style_label, 0),
'num_categories': FLAGS.num_styles,
'center': True,
'scale': True
}
# Dummy call to simplify the reuse logic
model.transform(inputs,
reuse=False,
normalizer_params=_create_normalizer_params(labels[0]))
def _style_sweep(inputs):
"""Transfers all styles onto the input one at a time."""
inputs = tf.expand_dims(inputs, 0)
stylized_inputs = [
model.transform(
inputs,
reuse=True,
normalizer_params=_create_normalizer_params(style_label))
for _, style_label in enumerate(labels)
]
return tf.concat_v2([inputs] + stylized_inputs, 0)
if FLAGS.style_grid:
style_row = tf.concat_v2([
tf.ones([1, FLAGS.image_size, FLAGS.image_size, 3]),
style_images
], 0)
stylized_training_example = _style_sweep(inputs[0])
stylized_evaluation_images = [
_style_sweep(image) for image in tf.unstack(evaluation_images)
]
stylized_noise = _style_sweep(
tf.random_uniform([FLAGS.image_size, FLAGS.image_size, 3]))
stylized_style_images = [
_style_sweep(image) for image in tf.unstack(style_images)
]
if FLAGS.style_crossover:
grid = tf.concat_v2(
[style_row, stylized_training_example, stylized_noise] +
stylized_evaluation_images + stylized_style_images, 0)
else:
grid = tf.concat_v2(
[style_row, stylized_training_example, stylized_noise] +
stylized_evaluation_images, 0)
tf.summary.image(
'Style Grid',
tf.cast(
image_utils.form_image_grid(grid, ([
3 + evaluation_images.get_shape().as_list()[0] +
FLAGS.num_styles, 1 + FLAGS.num_styles
] if FLAGS.style_crossover else [
3 + evaluation_images.get_shape().as_list()[0], 1 +
FLAGS.num_styles
]), [FLAGS.image_size, FLAGS.image_size], 3) * 255.0,
tf.uint8))
if FLAGS.learning_curves:
metrics = {}
for i, label in enumerate(labels):
gram_matrices = dict([
(key, value[i:i + 1])
for key, value in style_gram_matrices.iteritems()
])
stylized_inputs = model.transform(
inputs,
reuse=True,
normalizer_params=_create_normalizer_params(label))
_, loss_dict = learning.total_loss(inputs,
stylized_inputs,
gram_matrices,
content_weights,
style_weights,
reuse=i > 0)
for key, value in loss_dict.iteritems():
metrics['{}_style_{}'.format(
key, i)] = slim.metrics.streaming_mean(value)
names_values, names_updates = slim.metrics.aggregate_metric_map(
metrics)
for name, value in names_values.iteritems():
summary_op = tf.summary.scalar(name, value, [])
print_op = tf.Print(summary_op, [value], name)
tf.add_to_collection(tf.GraphKeys.SUMMARIES, print_op)
eval_op = names_updates.values()
num_evals = FLAGS.num_evals
else:
eval_op = None
num_evals = 1
slim.evaluation.evaluation_loop(
master=FLAGS.master,
checkpoint_dir=os.path.expanduser(FLAGS.train_dir),
logdir=os.path.expanduser(FLAGS.eval_dir),
eval_op=eval_op,
num_evals=num_evals,
eval_interval_secs=FLAGS.eval_interval_secs)
def main(_):
with tf.Graph().as_default():
# Create inputs in [0, 1], as expected by vgg_16.
inputs, _ = image_utils.imagenet_inputs(
FLAGS.batch_size, FLAGS.image_size)
evaluation_images = image_utils.load_evaluation_images(FLAGS.image_size)
# Process style and weight flags
if FLAGS.style_coefficients is None:
style_coefficients = [1.0 for _ in range(FLAGS.num_styles)]
else:
style_coefficients = ast.literal_eval(FLAGS.style_coefficients)
if len(style_coefficients) != FLAGS.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)
# Load style images.
style_images, labels, style_gram_matrices = image_utils.style_image_inputs(
os.path.expanduser(FLAGS.style_dataset_file),
batch_size=FLAGS.num_styles, image_size=FLAGS.image_size,
square_crop=True, shuffle=False)
labels = tf.unstack(labels)
def _create_normalizer_params(style_label):
"""Creates normalizer parameters from a style label."""
return {'labels': tf.expand_dims(style_label, 0),
'num_categories': FLAGS.num_styles,
'center': True,
'scale': True}
# Dummy call to simplify the reuse logic
model.transform(inputs, reuse=False,
normalizer_params=_create_normalizer_params(labels[0]))
def _style_sweep(inputs):
"""Transfers all styles onto the input one at a time."""
inputs = tf.expand_dims(inputs, 0)
stylized_inputs = [
model.transform(
inputs,
reuse=True,
normalizer_params=_create_normalizer_params(style_label))
for _, style_label in enumerate(labels)]
return tf.concat([inputs] + stylized_inputs, 0)
if FLAGS.style_grid:
style_row = tf.concat(
[tf.ones([1, FLAGS.image_size, FLAGS.image_size, 3]), style_images],
0)
stylized_training_example = _style_sweep(inputs[0])
stylized_evaluation_images = [
_style_sweep(image) for image in tf.unstack(evaluation_images)]
stylized_noise = _style_sweep(
tf.random_uniform([FLAGS.image_size, FLAGS.image_size, 3]))
stylized_style_images = [
_style_sweep(image) for image in tf.unstack(style_images)]
if FLAGS.style_crossover:
grid = tf.concat(
[style_row, stylized_training_example, stylized_noise] +
stylized_evaluation_images + stylized_style_images,
0)
else:
grid = tf.concat(
[style_row, stylized_training_example, stylized_noise] +
stylized_evaluation_images,
0)
if FLAGS.style_crossover:
grid_shape = [
3 + evaluation_images.get_shape().as_list()[0] + FLAGS.num_styles,
1 + FLAGS.num_styles]
else:
grid_shape = [
3 + evaluation_images.get_shape().as_list()[0],
1 + FLAGS.num_styles]
tf.summary.image(
'Style Grid',
tf.cast(
image_utils.form_image_grid(
grid,
grid_shape,
[FLAGS.image_size, FLAGS.image_size],
3) * 255.0,
tf.uint8))
if FLAGS.learning_curves:
metrics = {}
for i, label in enumerate(labels):
gram_matrices = dict(
(key, value[i: i + 1])
for key, value in style_gram_matrices.items())
stylized_inputs = model.transform(
inputs,
reuse=True,
normalizer_params=_create_normalizer_params(label))
_, loss_dict = learning.total_loss(
inputs, stylized_inputs, gram_matrices, content_weights,
style_weights, reuse=i > 0)
for key, value in loss_dict.items():
metrics['{}_style_{}'.format(key, i)] = slim.metrics.streaming_mean(
value)
names_values, names_updates = slim.metrics.aggregate_metric_map(metrics)
for name, value in names_values.items():
summary_op = tf.summary.scalar(name, value, [])
print_op = tf.Print(summary_op, [value], name)
tf.add_to_collection(tf.GraphKeys.SUMMARIES, print_op)
eval_op = names_updates.values()
num_evals = FLAGS.num_evals
else:
eval_op = None
num_evals = 1
slim.evaluation.evaluation_loop(
master=FLAGS.master,
checkpoint_dir=os.path.expanduser(FLAGS.train_dir),
logdir=os.path.expanduser(FLAGS.eval_dir),
eval_op=eval_op,
num_evals=num_evals,
eval_interval_secs=FLAGS.eval_interval_secs)
def main(_):
with tf.Graph().as_default():
# Create inputs in [0, 1], as expected by vgg_16.
inputs, _ = image_utils.imagenet_inputs(FLAGS.batch_size,
FLAGS.image_size)
evaluation_images = image_utils.load_evaluation_images(
FLAGS.image_size)
# Load style images.
style_images, labels, style_gram_matrices = image_utils.style_image_inputs(
os.path.expanduser(FLAGS.style_dataset_file),
batch_size=FLAGS.num_styles,
image_size=FLAGS.image_size,
square_crop=True,
shuffle=False)
labels = tf.unstack(labels)
def _create_normalizer_params(style_label):
"""Creates normalizer parameters from a style label."""
return {
'labels': tf.expand_dims(style_label, 0),
'num_categories': FLAGS.num_styles,
'center': True,
'scale': True
}
# Dummy call to simplify the reuse logic
model.transform(inputs,
alpha=FLAGS.alpha,
reuse=False,
normalizer_params=_create_normalizer_params(labels[0]))
def _style_sweep(inputs):
"""Transfers all styles onto the input one at a time."""
inputs = tf.expand_dims(inputs, 0)
stylized_inputs = []
for _, style_label in enumerate(labels):
stylized_input = model.transform(
inputs,
alpha=FLAGS.alpha,
reuse=True,
normalizer_params=_create_normalizer_params(style_label))
stylized_inputs.append(stylized_input)
return tf.concat([inputs] + stylized_inputs, 0)
style_row = tf.concat([
tf.ones([1, FLAGS.image_size, FLAGS.image_size, 3]), style_images
], 0)
stylized_training_example = _style_sweep(inputs[0])
stylized_evaluation_images = [
_style_sweep(image) for image in tf.unstack(evaluation_images)
]
stylized_noise = _style_sweep(
tf.random_uniform([FLAGS.image_size, FLAGS.image_size, 3]))
stylized_style_images = [
_style_sweep(image) for image in tf.unstack(style_images)
]
if FLAGS.style_crossover:
grid = tf.concat(
[style_row, stylized_training_example, stylized_noise] +
stylized_evaluation_images + stylized_style_images, 0)
else:
grid = tf.concat(
[style_row, stylized_training_example, stylized_noise] +
stylized_evaluation_images, 0)
if FLAGS.style_crossover:
grid_shape = [
3 + evaluation_images.get_shape().as_list()[0] +
FLAGS.num_styles, 1 + FLAGS.num_styles
]
else:
grid_shape = [
3 + evaluation_images.get_shape().as_list()[0],
1 + FLAGS.num_styles
]
style_grid = tf.cast(
image_utils.form_image_grid(
grid, grid_shape, [FLAGS.image_size, FLAGS.image_size], 3) *
255.0, tf.uint8)
sess = tf.Session()
with sess.as_default():
np_array = tf.squeeze(style_grid).eval()
im = Image.fromarray(np_array)
im.save('matrix.png')
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.iteritems()])
# Define the model
stylized_inputs = model.transform(
inputs,
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.iteritems():
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.
# TODO(iansimon): This is ugly, but assign_from_checkpoint_fn doesn't
# exist yet.
saver_vgg = tf.train.Saver(slim.get_variables('vgg_16'))
def init_fn_vgg(session):
saver_vgg.restore(session, vgg.checkpoint_file())
# Function to restore N-styles parameters.
# TODO(iansimon): This is ugly, but assign_from_checkpoint_fn doesn't
# exist yet.
saver_n_styles = tf.train.Saver(other_vars)
def init_fn_n_styles(session):
saver_n_styles.restore(session, os.path.expanduser(FLAGS.checkpoint))
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)
# 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,
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):
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)):
# Load 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, \
true_loss, \
_, \
bottleneck_feat = build_mobilenet_model.build_mobilenet_model(
content_inputs_,
style_inputs_,
mobilenet_trainable=True,
style_params_trainable=False,
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=FLAGS.total_variation_weight,
)
_, inception_bottleneck_feat = build_model.style_prediction(
style_inputs_,
[],
[],
is_training=False,
trainable=False,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
reuse=None,
)
print('PRINTING TRAINABLE VARIABLES')
for x in tf.trainable_variables():
print(x)
mse_loss = tf.losses.mean_squared_error(
inception_bottleneck_feat, bottleneck_feat)
total_loss = mse_loss
if FLAGS.use_true_loss:
true_loss = FLAGS.true_loss_weight*true_loss
total_loss += true_loss
if FLAGS.use_true_loss:
tf.summary.scalar('mse', mse_loss)
tf.summary.scalar('true_loss', true_loss)
tf.summary.scalar('total_loss', total_loss)
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)
mobilenet_variables_to_restore = slim.get_variables_to_restore(
include=['MobilenetV2'],
exclude=['global_step'])
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
)
init_fn = slim.assign_from_checkpoint_fn(
FLAGS.initial_checkpoint,
slim.get_variables_to_restore(
exclude=['MobilenetV2', 'mobilenet_conv', 'global_step']))
init_pretrained_mobilenet = slim.assign_from_checkpoint_fn(
FLAGS.mobilenet_checkpoint, mobilenet_variables_to_restore)
def init_sub_networks(session):
init_fn(session)
init_pretrained_mobilenet(session)
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)
