text: # Copyright 2016 The TensorFlow Authors All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  """Training for Domain Separation Networks (DSNs)."""
  from __future__ import division
  
  import dsn
  import tensorflow as tf
  from domain_adaptation.datasets import dataset_factory
  
  slim = tf.contrib.slim
  FLAGS = tf.app.flags.FLAGS
  
  tf.app.flags.DEFINE_integer('batch_size', 32,
                              'The number of images in each batch.')
  
  tf.app.flags.DEFINE_string('source_dataset', 'pose_synthetic',
                             'Source dataset to train on.')
  
  tf.app.flags.DEFINE_string('target_dataset', 'pose_real',
                             'Target dataset to train on.')
  
  tf.app.flags.DEFINE_string('target_labeled_dataset', 'none',
                             'Target dataset to train on.')
  
  tf.app.flags.DEFINE_string('dataset_dir', None,
                             'The directory where the dataset files are stored.')
  
  tf.app.flags.DEFINE_string('master', '',
                             'BNS name of the TensorFlow master to use.')
  
  tf.app.flags.DEFINE_string('train_log_dir', '/tmp/da/',
                             'Directory where to write event logs.')
  
  tf.app.flags.DEFINE_string(
      'layers_to_regularize', 'fc3',
      'Comma-separated list of layer names to use MMD regularization on.')
  
  tf.app.flags.DEFINE_float('learning_rate', .01, 'The learning rate')
  
  tf.app.flags.DEFINE_float('alpha_weight', 1e-6,
                            'The coefficient for scaling the reconstruction '
                            'loss.')
  
  tf.app.flags.DEFINE_float(
      'beta_weight', 1e-6,
      'The coefficient for scaling the private/shared difference loss.')
  
  tf.app.flags.DEFINE_float(
      'gamma_weight', 1e-6,
      'The coefficient for scaling the shared encoding similarity loss.')
  
  tf.app.flags.DEFINE_float('pose_weight', 0.125,
                            'The coefficient for scaling the pose loss.')
  
  tf.app.flags.DEFINE_float(
      'weight_decay', 1e-6,
      'The coefficient for the L2 regularization applied for all weights.')
  
  tf.app.flags.DEFINE_integer(
      'save_summaries_secs', 60,
      'The frequency with which summaries are saved, in seconds.')
  
  tf.app.flags.DEFINE_integer(
      'save_interval_secs', 60,
      'The frequency with which the model is saved, in seconds.')
  
  tf.app.flags.DEFINE_integer(
      'max_number_of_steps', None,
      'The maximum number of gradient steps. Use None to train indefinitely.')
  
  tf.app.flags.DEFINE_integer(
      'domain_separation_startpoint', 1,
      'The global step to add the domain separation losses.')
  
  tf.app.flags.DEFINE_integer(
      'bipartite_assignment_top_k', 3,
      'The number of top-k matches to use in bipartite matching adaptation.')
  
  tf.app.flags.DEFINE_float('decay_rate', 0.95, 'Learning rate decay factor.')
  
  tf.app.flags.DEFINE_integer('decay_steps', 20000, 'Learning rate decay steps.')
  
  tf.app.flags.DEFINE_float('momentum', 0.9, 'The momentum value.')
  
  tf.app.flags.DEFINE_bool('use_separation', False,
                           'Use our domain separation model.')
  
  tf.app.flags.DEFINE_bool('use_logging', False, 'Debugging messages.')
  
  tf.app.flags.DEFINE_integer(
      'ps_tasks', 0,
      'The number of parameter servers. If the value is 0, then the parameters '
      'are handled locally by the worker.')
  
  tf.app.flags.DEFINE_integer(
      'num_readers', 4,
      'The number of parallel readers that read data from the dataset.')
  
  tf.app.flags.DEFINE_integer('num_preprocessing_threads', 4,
                              'The number of threads used to create the batches.')
  
  tf.app.flags.DEFINE_integer(
      'task', 0,
      'The Task ID. This value is used when training with multiple workers to '
      'identify each worker.')
  
  tf.app.flags.DEFINE_string('decoder_name', 'small_decoder',
                             'The decoder to use.')
  tf.app.flags.DEFINE_string('encoder_name', 'default_encoder',
                             'The encoder to use.')
  
  ################################################################################
  # Flags that control the architecture and losses
  ################################################################################
  tf.app.flags.DEFINE_string(
      'similarity_loss', 'grl',
      'The method to use for encouraging the common encoder codes to be '
      'similar, one of "grl", "mmd", "corr".')
  
  tf.app.flags.DEFINE_string('recon_loss_name', 'sum_of_pairwise_squares',
                             'The name of the reconstruction loss.')
  
  tf.app.flags.DEFINE_string('basic_tower', 'pose_mini',
                             'The basic tower building block.')
  
  def provide_batch_fn():
    """ The provide_batch function to use. """
    return dataset_factory.provide_batch
  
  def main(_):
    model_params = {
        'use_separation': FLAGS.use_separation,
        'domain_separation_startpoint': FLAGS.domain_separation_startpoint,
        'layers_to_regularize': FLAGS.layers_to_regularize,
        'alpha_weight': FLAGS.alpha_weight,
        'beta_weight': FLAGS.beta_weight,
        'gamma_weight': FLAGS.gamma_weight,
        'pose_weight': FLAGS.pose_weight,
        'recon_loss_name': FLAGS.recon_loss_name,
        'decoder_name': FLAGS.decoder_name,
        'encoder_name': FLAGS.encoder_name,
        'weight_decay': FLAGS.weight_decay,
        'batch_size': FLAGS.batch_size,
        'use_logging': FLAGS.use_logging,
        'ps_tasks': FLAGS.ps_tasks,
        'task': FLAGS.task,
    }
    g = tf.Graph()
    with g.as_default():
      with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
        # Load the data.
        source_images, source_labels = provide_batch_fn()(
            FLAGS.source_dataset, 'train', FLAGS.dataset_dir, FLAGS.num_readers,
            FLAGS.batch_size, FLAGS.num_preprocessing_threads)
        target_images, target_labels = provide_batch_fn()(
            FLAGS.target_dataset, 'train', FLAGS.dataset_dir, FLAGS.num_readers,
            FLAGS.batch_size, FLAGS.num_preprocessing_threads)
  
        # In the unsupervised case all the samples in the labeled
        # domain are from the source domain.
        domain_selection_mask = tf.fill((source_images.get_shape().as_list()[0],),
                                        True)
  
        # When using the semisupervised model we include labeled target data in
        # the source labelled data.
        if FLAGS.target_labeled_dataset != 'none':
          # 1000 is the maximum number of labelled target samples that exists in
          # the datasets.
          target_semi_images, target_semi_labels = provide_batch_fn()(
              FLAGS.target_labeled_dataset, 'train', FLAGS.batch_size)
  
          # Calculate the proportion of source domain samples in the semi-
          # supervised setting, so that the proportion is set accordingly in the
          # batches.
          proportion = float(source_labels['num_train_samples']) / (
              source_labels['num_train_samples'] +
              target_semi_labels['num_train_samples'])
  
          rnd_tensor = tf.random_uniform(
              (target_semi_images.get_shape().as_list()[0],))
  
          domain_selection_mask = rnd_tensor < proportion
          source_images = tf.where(domain_selection_mask, source_images,
                                   target_semi_images)
          source_class_labels = tf.where(domain_selection_mask,
                                         source_labels['classes'],
                                         target_semi_labels['classes'])
  
          if 'quaternions' in source_labels:
            source_pose_labels = tf.where(domain_selection_mask,
                                          source_labels['quaternions'],
                                          target_semi_labels['quaternions'])
            (source_images, source_class_labels, source_pose_labels,
             domain_selection_mask) = tf.train.shuffle_batch(
                 [
                     source_images, source_class_labels, source_pose_labels,
                     domain_selection_mask
                 ],
                 FLAGS.batch_size,
                 50000,
                 5000,
                 num_threads=1,
                 enqueue_many=True)
  
          else:
            (source_images, source_class_labels,
             domain_selection_mask) = tf.train.shuffle_batch(
                 [source_images, source_class_labels, domain_selection_mask],
                 FLAGS.batch_size,
                 50000,
                 5000,
                 num_threads=1,
                 enqueue_many=True)
          source_labels = {}
          source_labels['classes'] = source_class_labels
          if 'quaternions' in source_labels:
            source_labels['quaternions'] = source_pose_labels
  
        slim.get_or_create_global_step()
        tf.summary.image('source_images', source_images, max_outputs=3)
        tf.summary.image('target_images', target_images, max_outputs=3)
  
        dsn.create_model(
            source_images,
            source_labels,
            domain_selection_mask,
            target_images,
            target_labels,
            FLAGS.similarity_loss,
            model_params,
            basic_tower_name=FLAGS.basic_tower)
  
        # Configure the optimization scheme:
        learning_rate = tf.train.exponential_decay(
            FLAGS.learning_rate,
            slim.get_or_create_global_step(),
            FLAGS.decay_steps,
            FLAGS.decay_rate,
            staircase=True,
            name='learning_rate')
  
        tf.summary.scalar('learning_rate', learning_rate)
        tf.summary.scalar('total_loss', tf.losses.get_total_loss())
  
        opt = tf.train.MomentumOptimizer(learning_rate, FLAGS.momentum)
        tf.logging.set_verbosity(tf.logging.INFO)
        # Run training.
        loss_tensor = slim.learning.create_train_op(
            slim.losses.get_total_loss(),
            opt,
            summarize_gradients=True,
            colocate_gradients_with_ops=True)
        slim.learning.train(
            train_op=loss_tensor,
            logdir=FLAGS.train_log_dir,
            master=FLAGS.master,
            is_chief=FLAGS.task == 0,
            number_of_steps=FLAGS.max_number_of_steps,
            save_summaries_secs=FLAGS.save_summaries_secs,
            save_interval_secs=FLAGS.save_interval_secs)
  
  
  if __name__ == '__main__':
    tf.app.run()
  