Run a network on the PTB training set, checkpointing the weights.

def train(logdir):
  """Run a network on the PTB training set, checkpointing the weights."""

  ptb_train = PTB(
      name="ptb_train",
      subset="train",
      seq_len=FLAGS.unroll_steps,
      batch_size=FLAGS.batch_size)

  # Connect to training set.
  data_ops = ptb_train()
  embed_layer, rnn_core, output_linear = build_modules(
      is_training=True, vocab_size=ptb_train.vocab_size)
  prediction_logits, zero_state_op = build_logits(
      data_ops, embed_layer, rnn_core, output_linear, name_prefix="train")
  data_loss = build_loss(prediction_logits, data_ops.sparse_target)

  # Add the KL cost.
  total_kl_cost = bbb.get_total_kl_cost()
  num_dataset_elements = FLAGS.batch_size * ptb_train.num_batches
  scaled_kl_cost = total_kl_cost / num_dataset_elements
  total_loss = tf.add(scaled_kl_cost, data_loss)

  # Optimize as usual.
  global_step = tf.get_variable(
      "num_weight_updates",
      initializer=tf.constant(0, dtype=tf.int32, shape=()),
      collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP])

  learning_rate = tf.get_variable(
      "lr", initializer=tf.constant(FLAGS.lr_start, shape=(), dtype=tf.float32))
  learning_rate_update = learning_rate.assign(learning_rate * FLAGS.lr_decay)
  optimizer = tf.train.GradientDescentOptimizer(
      learning_rate=learning_rate)
  optimizer = GlobalNormClippingOptimizer(optimizer, clip_norm=5.0)

  with tf.control_dependencies([optimizer.minimize(total_loss)]):
    global_step_and_train = global_step.assign_add(1)

  # Connect to valid set.
  ptb_valid = PTB(
      name="ptb_valid",
      subset="valid",
      seq_len=FLAGS.unroll_steps,
      batch_size=FLAGS.batch_size)
  valid_data_ops = ptb_valid()
  valid_logits, zero_valid_state = build_logits(
      valid_data_ops, embed_layer, rnn_core, output_linear, name_prefix="valid")
  valid_loss = build_loss(valid_logits, valid_data_ops.sparse_target)

  # Compute metrics for the sake of monitoring training.
  predictions = tf.cast(
      tf.argmax(prediction_logits, axis=-1), tf.int32, name="pred")
  correct_prediction_mask = tf.cast(
      tf.equal(predictions, data_ops.sparse_target), tf.int32)
  accuracy = tf.reduce_mean(
      tf.cast(correct_prediction_mask, tf.float32), name="acc")
  error_rate = tf.subtract(1.0, accuracy, name="err")
  label_probs = tf.nn.softmax(prediction_logits, dim=-1)
  predictive_entropy = tf.reduce_mean(
      label_probs * tf.log(label_probs + 1e-12) * -1.0)

  # Create tf.summary ops.
  log_ops_to_run = {
      "scalar": collections.OrderedDict([
          ("task_loss", data_loss),
          ("train_err_rate", error_rate),
          ("pred_entropy", predictive_entropy),
          ("learning_rate", learning_rate),
          ("elbo_loss", total_loss),
          ("kl_cost", total_kl_cost),
          ("scaled_kl_cost", scaled_kl_cost),
      ]),
      "text": collections.OrderedDict([
          ("labels", ptb_train.to_string_tensor(data_ops.sparse_target)),
          ("predictions", ptb_train.to_string_tensor(predictions))
      ])
  }

  for name, tensor in log_ops_to_run["scalar"].items():
    tf.summary.scalar(os.path.join("train", name), tensor)

  # The remaining logic runs the training loop and logging.
  summary_writer = tf.summary.FileWriterCache.get(logdir=logdir)
  tf.logging.info(
      "Beginning training for {} epochs, each with {} batches.".format(
          FLAGS.num_training_epochs, ptb_train.num_batches))
  with tf.train.MonitoredTrainingSession(
      is_chief=True, checkpoint_dir=logdir, save_summaries_secs=10) as sess:
    num_updates_v = _run_session_with_no_hooks(sess, global_step)
    epoch_idx_start, step_idx_start = divmod(
        num_updates_v, ptb_train.num_batches)
    tf.logging.info("On start, epoch: {}\t step: {}".format(
        epoch_idx_start, step_idx_start))
    for epoch_idx in range(epoch_idx_start, FLAGS.num_training_epochs):
      tf.logging.info("Beginning Epoch {}/{}".format(
          epoch_idx, FLAGS.num_training_epochs))
      tf.logging.info(
          ("Beginning by evaluating on the validation set, which has "
           "{} batches.".format(ptb_valid.num_batches)))
      valid_cost = 0
      valid_steps = 0
      _run_session_with_no_hooks(sess, zero_valid_state)
      for _ in range(ptb_valid.num_batches):
        valid_cost_v, num_updates_v = _run_session_with_no_hooks(
            sess, [valid_loss, global_step])
        valid_cost += valid_cost_v
        valid_steps += FLAGS.unroll_steps
      tf.logging.info("Validation set perplexity: {}".format(
          np.exp(valid_cost / valid_steps)))
      summary = tf.summary.Summary()
      summary.value.add(
          tag="valid/word_level_perplexity",
          simple_value=np.exp(valid_cost / valid_steps))
      summary_writer.add_summary(summary, num_updates_v)

      # Run a training epoch.
      epoch_cost = 0
      epoch_steps = 0
      for batch_idx in range(step_idx_start, ptb_train.num_batches):
        scalars_res, num_updates_v = sess.run(
            [log_ops_to_run["scalar"], global_step_and_train])
        epoch_cost += scalars_res["task_loss"]
        epoch_steps += FLAGS.unroll_steps
        if (batch_idx - 1) % FLAGS.print_every_batches == 0:
          summary = tf.summary.Summary()
          summary.value.add(
              tag="train/word_level_perplexity",
              simple_value=np.exp(epoch_cost / epoch_steps))
          summary_writer.add_summary(summary, num_updates_v)
          scalars_res, strings_res = _run_session_with_no_hooks(
              sess, [log_ops_to_run["scalar"], log_ops_to_run["text"]])
          tf.logging.info("Num weight updates: {}".format(num_updates_v))
          for name, result in scalars_res.items():
            tf.logging.info("{}: {}".format(name, result))
          for name, result in strings_res.items():
            tf.logging.info("{}: {}".format(name, result))

      word_level_perplexity = np.exp(epoch_cost / epoch_steps)
      tf.logging.info(
          "Train Perplexity after Epoch {}: {}".format(
              epoch_idx, word_level_perplexity))

      end_of_epoch_fetches = [zero_state_op]
      if epoch_idx >= FLAGS.high_lr_epochs:
        end_of_epoch_fetches.append(learning_rate_update)
      _run_session_with_no_hooks(sess, end_of_epoch_fetches)

  tf.logging.info("Done training. Thanks for your time.")