(FLAGS, features, labels, mems, is_training)
return total_loss, new_mems, monitor_dict
Pretraining loss with two-stream attention Transformer-XL.
Pretraining loss with two-stream attention Transformer-XL.
Pretraining loss with two - stream attention Transformer - XL .
def two_stream_loss(FLAGS, features, labels, mems, is_training):
  """Pretraining loss with two-stream attention Transformer-XL."""

  #### Unpack input
  mem_name = "mems"
  mems = mems.get(mem_name, None)

  inp_k = tf.transpose(features["input_k"], [1, 0])
  inp_q = tf.transpose(features["input_q"], [1, 0])

  seg_id = tf.transpose(features["seg_id"], [1, 0])

  inp_mask = None
  perm_mask = tf.transpose(features["perm_mask"], [1, 2, 0])

  if FLAGS.num_predict is not None:
    # [num_predict x tgt_len x bsz]
    target_mapping = tf.transpose(features["target_mapping"], [1, 2, 0])
  else:
    target_mapping = None

  # target for LM loss
  tgt = tf.transpose(features["target"], [1, 0])

  # target mask for LM loss
  tgt_mask = tf.transpose(features["target_mask"], [1, 0])

  # construct xlnet config and save to model_dir
  xlnet_config = xlnet.XLNetConfig(FLAGS=FLAGS)
  xlnet_config.to_json(os.path.join(FLAGS.model_dir, "config.json"))

  # construct run config from FLAGS
  run_config = xlnet.create_run_config(is_training, False, FLAGS)

  xlnet_model = xlnet.XLNetModel(
      xlnet_config=xlnet_config,
      run_config=run_config,
      input_ids=inp_k,
      seg_ids=seg_id,
      input_mask=inp_mask,
      mems=mems,
      perm_mask=perm_mask,
      target_mapping=target_mapping,
      inp_q=inp_q)

  output = xlnet_model.get_sequence_output()
  new_mems = {mem_name: xlnet_model.get_new_memory()}
  lookup_table = xlnet_model.get_embedding_table()

  initializer = xlnet_model.get_initializer()

  with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
    # LM loss
    lm_loss = modeling.lm_loss(
        hidden=output,
        target=tgt,
        n_token=xlnet_config.n_token,
        d_model=xlnet_config.d_model,
        initializer=initializer,
        lookup_table=lookup_table,
        tie_weight=True,
        bi_data=run_config.bi_data,
        use_tpu=run_config.use_tpu)

  #### Quantity to monitor
  monitor_dict = {}

  if FLAGS.use_bfloat16:
    tgt_mask = tf.cast(tgt_mask, tf.float32)
    lm_loss = tf.cast(lm_loss, tf.float32)

  total_loss = tf.reduce_sum(lm_loss * tgt_mask) / tf.reduce_sum(tgt_mask)
  monitor_dict["total_loss"] = total_loss

  return total_loss, new_mems, monitor_dict
def two_stream_loss ( FLAGS , features , labels , mems , is_training ) : #### Unpack input mem_name = "mems" mems = mems . get ( mem_name , None ) inp_k = tf . transpose ( features [ "input_k" ] , [ 1 , 0 ] ) inp_q = tf . transpose ( features [ "input_q" ] , [ 1 , 0 ] ) seg_id = tf . transpose ( features [ "seg_id" ] , [ 1 , 0 ] ) inp_mask = None perm_mask = tf . transpose ( features [ "perm_mask" ] , [ 1 , 2 , 0 ] ) if FLAGS . num_predict is not None : # [num_predict x tgt_len x bsz] target_mapping = tf . transpose ( features [ "target_mapping" ] , [ 1 , 2 , 0 ] ) else : target_mapping = None # target for LM loss tgt = tf . transpose ( features [ "target" ] , [ 1 , 0 ] ) # target mask for LM loss tgt_mask = tf . transpose ( features [ "target_mask" ] , [ 1 , 0 ] ) # construct xlnet config and save to model_dir xlnet_config = xlnet . XLNetConfig ( FLAGS = FLAGS ) xlnet_config . to_json ( os . path . join ( FLAGS . model_dir , "config.json" ) ) # construct run config from FLAGS run_config = xlnet . create_run_config ( is_training , False , FLAGS ) xlnet_model = xlnet . XLNetModel ( xlnet_config = xlnet_config , run_config = run_config , input_ids = inp_k , seg_ids = seg_id , input_mask = inp_mask , mems = mems , perm_mask = perm_mask , target_mapping = target_mapping , inp_q = inp_q ) output = xlnet_model . get_sequence_output ( ) new_mems = { mem_name : xlnet_model . get_new_memory ( ) } lookup_table = xlnet_model . get_embedding_table ( ) initializer = xlnet_model . get_initializer ( ) with tf . variable_scope ( "model" , reuse = tf . AUTO_REUSE ) : # LM loss lm_loss = modeling . lm_loss ( hidden = output , target = tgt , n_token = xlnet_config . n_token , d_model = xlnet_config . d_model , initializer = initializer , lookup_table = lookup_table , tie_weight = True , bi_data = run_config . bi_data , use_tpu = run_config . use_tpu ) #### Quantity to monitor monitor_dict = { } if FLAGS . use_bfloat16 : tgt_mask = tf . cast ( tgt_mask , tf . float32 ) lm_loss = tf . cast ( lm_loss , tf . float32 ) total_loss = tf . reduce_sum ( lm_loss * tgt_mask ) / tf . reduce_sum ( tgt_mask ) monitor_dict [ "total_loss" ] = total_loss return total_loss , new_mems , monitor_dict