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