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 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)