text: # coding=utf-8
  # Copyright 2020 The Google Research Authors.
  #
  # 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.
  
  """Train model."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl import app
  from absl import flags
  import numpy as np
  import tensorflow.compat.v1 as tf
  
  from large_margin import margin_loss
  from large_margin.mnist import data_provider as mnist
  from large_margin.mnist import mnist_config
  from large_margin.mnist import mnist_model
  
  flags.DEFINE_integer("num_epochs", 200, "Number of training epochs.")
  flags.DEFINE_string("master", "local",
                      "BNS name of the TensorFlow master to use.")
  flags.DEFINE_integer("task", 0, "Task id of the replica running the training.")
  flags.DEFINE_integer("ps_tasks", 0,
                       "Number of tasks in the ps job. If 0 no ps job is used.")
  flags.DEFINE_integer("num_replicas", 0, "Number of workers in the job.")
  flags.DEFINE_string("checkpoint_dir", "/tmp/large_margin/train/",
                      "Results directory.")
  flags.DEFINE_float("momentum", 0.0,
                     "Momentum value if used momentum optimizer.")
  flags.DEFINE_integer("decay_steps", 2000,
                       "Number of steps to decay learning rate")
  flags.DEFINE_float("decay_rate", 0.96, "Rate of decay")
  flags.DEFINE_integer("batch_size", 256, "Training batch size.")
  flags.DEFINE_integer("log_every_steps", 50,
                       "Saving logging frequency in optimization steps.")
  flags.DEFINE_integer("save_checkpoint_secs", 2 * 60,
                       "Saving checkpoints frequency in secs.")
  flags.DEFINE_integer("save_summaries_secs", 2 * 60,
                       "Saving summaries frequency in secs.")
  flags.DEFINE_enum("experiment_type", "mnist", ["mnist"], "Experiment type.")
  flags.DEFINE_integer("startup_delay_steps", 15,
                       "Number of training steps between replicas startup.")
  flags.DEFINE_string("data_dir", "", "Data directory.")
  
  FLAGS = flags.FLAGS
  
  
  def train():
    """Training function."""
    is_chief = (FLAGS.task == 0)
    g = tf.Graph()
    with g.as_default():
      with tf.device(tf.train.replica_device_setter(ps_tasks=FLAGS.ps_tasks)):
        if FLAGS.experiment_type == "mnist":
          config = mnist_config.ConfigDict()
          dataset = mnist.MNIST(
              data_dir=FLAGS.data_dir,
              subset="train",
              batch_size=FLAGS.batch_size,
              is_training=True)
          model = mnist_model.MNISTNetwork(config)
          layers_names = [
              "conv_layer%d" % i
              for i in range(len(config.filter_sizes_conv_layers))
          ]
  
        images, labels, num_examples, num_classes = (dataset.images,
                                                     dataset.labels,
                                                     dataset.num_examples,
                                                     dataset.num_classes)
        tf.summary.image("images", images)
  
        # Build model.
        logits, endpoints = model(images, is_training=True)
        layers_list = [images] + [endpoints[name] for name in layers_names]
  
        # Define losses.
        l2_loss_wt = config.l2_loss_wt
        xent_loss_wt = config.xent_loss_wt
        margin_loss_wt = config.margin_loss_wt
        gamma = config.gamma
        alpha = config.alpha
        top_k = config.top_k
        dist_norm = config.dist_norm
        with tf.name_scope("losses"):
          xent_loss = tf.reduce_mean(
              tf.nn.sparse_softmax_cross_entropy_with_logits(
                  logits=logits, labels=labels))
          margin = margin_loss.large_margin(
              logits=logits,
              one_hot_labels=tf.one_hot(labels, num_classes),
              layers_list=layers_list,
              gamma=gamma,
              alpha_factor=alpha,
              top_k=top_k,
              dist_norm=dist_norm,
              epsilon=1e-6,
              layers_weights=[
                  np.prod(layer.get_shape().as_list()[1:])
                  for layer in layers_list] if np.isinf(dist_norm) else None
              )
  
          l2_loss = 0.
          for v in tf.trainable_variables():
            tf.logging.info(v)
            l2_loss += tf.nn.l2_loss(v)
  
          total_loss = 0
          if xent_loss_wt > 0:
            total_loss += xent_loss_wt * xent_loss
          if margin_loss_wt > 0:
            total_loss += margin_loss_wt * margin
          if l2_loss_wt > 0:
            total_loss += l2_loss_wt * l2_loss
  
          tf.summary.scalar("xent_loss", xent_loss)
          tf.summary.scalar("margin_loss", margin)
          tf.summary.scalar("l2_loss", l2_loss)
          tf.summary.scalar("total_loss", total_loss)
  
        # Build optimizer.
        init_lr = config.init_lr
        with tf.name_scope("optimizer"):
          global_step = tf.train.get_or_create_global_step()
          if FLAGS.num_replicas > 1:
            num_batches_per_epoch = num_examples // (
                FLAGS.batch_size * FLAGS.num_replicas)
          else:
            num_batches_per_epoch = num_examples // FLAGS.batch_size
          max_iters = num_batches_per_epoch * FLAGS.num_epochs
  
          lr = tf.train.exponential_decay(init_lr,
                                          global_step,
                                          FLAGS.decay_steps,
                                          FLAGS.decay_rate,
                                          staircase=True,
                                          name="lr_schedule")
  
          tf.summary.scalar("learning_rate", lr)
  
          var_list = tf.trainable_variables()
          grad_vars = tf.gradients(total_loss, var_list)
          tf.summary.scalar(
              "grad_norm",
              tf.reduce_mean([tf.norm(grad_var) for grad_var in grad_vars]))
          grad_vars, _ = tf.clip_by_global_norm(grad_vars, 5.0)
  
          opt = tf.train.RMSPropOptimizer(lr, momentum=FLAGS.momentum,
                                          epsilon=1e-2)
          if FLAGS.num_replicas > 1:
            opt = tf.train.SyncReplicasOptimizer(
                opt,
                replicas_to_aggregate=FLAGS.num_replicas,
                total_num_replicas=FLAGS.num_replicas)
  
          update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
          with tf.control_dependencies(update_ops):
            opt_op = opt.apply_gradients(
                zip(grad_vars, var_list), global_step=global_step)
  
        # Compute accuracy.
        top1_op = tf.nn.in_top_k(logits, labels, 1)
        accuracy = tf.reduce_mean(tf.cast(top1_op, dtype=tf.float32))
        tf.summary.scalar("top1_accuracy", accuracy)
  
        # Prepare optimization.
        vars_to_save = tf.global_variables()
        saver = tf.train.Saver(var_list=vars_to_save, max_to_keep=5, sharded=True)
        merged_summary = tf.summary.merge_all()
  
        # Hooks for optimization.
        hooks = [tf.train.StopAtStepHook(last_step=max_iters)]
        if not is_chief:
          hooks.append(
              tf.train.GlobalStepWaiterHook(
                  FLAGS.task * FLAGS.startup_delay_steps))
  
        init_op = tf.global_variables_initializer()
        scaffold = tf.train.Scaffold(
            init_op=init_op, summary_op=merged_summary, saver=saver)
  
        # Run optimization.
        epoch = 0
        with tf.train.MonitoredTrainingSession(
            is_chief=is_chief,
            checkpoint_dir=FLAGS.checkpoint_dir,
            hooks=hooks,
            save_checkpoint_secs=FLAGS.save_checkpoint_secs,
            save_summaries_secs=FLAGS.save_summaries_secs,
            scaffold=scaffold) as sess:
          while not sess.should_stop():
            _, acc, i = sess.run((opt_op, accuracy, global_step))
            epoch = i // num_batches_per_epoch
            if (i % FLAGS.log_every_steps) == 0:
              tf.logging.info("global step %d: epoch %d:\n train accuracy %.3f" %
                              (i, epoch, acc))
  
  
  def main(argv):
    del argv  # Unused
    # Training.
    train()
  
  
  if __name__ == "__main__":
    app.run(main)
  