(bert_config, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps, use_tpu,
                     use_one_hot_embeddings)
return model_fn
Returns `model_fn` closure for TPUEstimator.
Returns `model_fn` closure for TPUEstimator.
Returns model_fn closure for TPUEstimator .
def model_fn_builder(bert_config, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps, use_tpu,
                     use_one_hot_embeddings):
  """Returns `model_fn` closure for TPUEstimator."""

  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
    """The `model_fn` for TPUEstimator."""

    tf.logging.info("*** Features ***")
    for name in sorted(features.keys()):
      tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))

    unique_ids = features["unique_ids"]
    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)

    (start_logits, end_logits) = create_model(
        bert_config=bert_config,
        is_training=is_training,
        input_ids=input_ids,
        input_mask=input_mask,
        segment_ids=segment_ids,
        use_one_hot_embeddings=use_one_hot_embeddings)

    tvars = tf.trainable_variables()

    initialized_variable_names = {}
    scaffold_fn = None
    if init_checkpoint:
      (assignment_map, initialized_variable_names
      ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
      if use_tpu:

        def tpu_scaffold():
          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
          return tf.train.Scaffold()

        scaffold_fn = tpu_scaffold
      else:
        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

    tf.logging.info("**** Trainable Variables ****")
    for var in tvars:
      init_string = ""
      if var.name in initialized_variable_names:
        init_string = ", *INIT_FROM_CKPT*"
      tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
                      init_string)

    output_spec = None
    if mode == tf.estimator.ModeKeys.TRAIN:
      seq_length = modeling.get_shape_list(input_ids)[1]

      def compute_loss(logits, positions):
        one_hot_positions = tf.one_hot(
            positions, depth=seq_length, dtype=tf.float32)
        log_probs = tf.nn.log_softmax(logits, axis=-1)
        loss = -tf.reduce_mean(
            tf.reduce_sum(one_hot_positions * log_probs, axis=-1))
        return loss

      start_positions = features["start_positions"]
      end_positions = features["end_positions"]

      start_loss = compute_loss(start_logits, start_positions)
      end_loss = compute_loss(end_logits, end_positions)

      total_loss = (start_loss + end_loss) / 2.0

      train_op = optimization.create_optimizer(
          total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)

      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          train_op=train_op,
          scaffold_fn=scaffold_fn)
    elif mode == tf.estimator.ModeKeys.PREDICT:
      predictions = {
          "unique_ids": unique_ids,
          "start_logits": start_logits,
          "end_logits": end_logits,
      }
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)
    else:
      raise ValueError(
          "Only TRAIN and PREDICT modes are supported: %s" % (mode))

    return output_spec

  return model_fn
def model_fn_builder ( bert_config , init_checkpoint , learning_rate , num_train_steps , num_warmup_steps , use_tpu , use_one_hot_embeddings ) : def model_fn ( features , labels , mode , params ) : # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" tf . logging . info ( "*** Features ***" ) for name in sorted ( features . keys ( ) ) : tf . logging . info ( "  name = %s, shape = %s" % ( name , features [ name ] . shape ) ) unique_ids = features [ "unique_ids" ] input_ids = features [ "input_ids" ] input_mask = features [ "input_mask" ] segment_ids = features [ "segment_ids" ] is_training = ( mode == tf . estimator . ModeKeys . TRAIN ) ( start_logits , end_logits ) = create_model ( bert_config = bert_config , is_training = is_training , input_ids = input_ids , input_mask = input_mask , segment_ids = segment_ids , use_one_hot_embeddings = use_one_hot_embeddings ) tvars = tf . trainable_variables ( ) initialized_variable_names = { } scaffold_fn = None if init_checkpoint : ( assignment_map , initialized_variable_names ) = modeling . get_assignment_map_from_checkpoint ( tvars , init_checkpoint ) if use_tpu : def tpu_scaffold ( ) : tf . train . init_from_checkpoint ( init_checkpoint , assignment_map ) return tf . train . Scaffold ( ) scaffold_fn = tpu_scaffold else : tf . train . init_from_checkpoint ( init_checkpoint , assignment_map ) tf . logging . info ( "**** Trainable Variables ****" ) for var in tvars : init_string = "" if var . name in initialized_variable_names : init_string = ", *INIT_FROM_CKPT*" tf . logging . info ( "  name = %s, shape = %s%s" , var . name , var . shape , init_string ) output_spec = None if mode == tf . estimator . ModeKeys . TRAIN : seq_length = modeling . get_shape_list ( input_ids ) [ 1 ] def compute_loss ( logits , positions ) : one_hot_positions = tf . one_hot ( positions , depth = seq_length , dtype = tf . float32 ) log_probs = tf . nn . log_softmax ( logits , axis = - 1 ) loss = - tf . reduce_mean ( tf . reduce_sum ( one_hot_positions * log_probs , axis = - 1 ) ) return loss start_positions = features [ "start_positions" ] end_positions = features [ "end_positions" ] start_loss = compute_loss ( start_logits , start_positions ) end_loss = compute_loss ( end_logits , end_positions ) total_loss = ( start_loss + end_loss ) / 2.0 train_op = optimization . create_optimizer ( total_loss , learning_rate , num_train_steps , num_warmup_steps , use_tpu ) output_spec = tf . contrib . tpu . TPUEstimatorSpec ( mode = mode , loss = total_loss , train_op = train_op , scaffold_fn = scaffold_fn ) elif mode == tf . estimator . ModeKeys . PREDICT : predictions = { "unique_ids" : unique_ids , "start_logits" : start_logits , "end_logits" : end_logits , } output_spec = tf . contrib . tpu . TPUEstimatorSpec ( mode = mode , predictions = predictions , scaffold_fn = scaffold_fn ) else : raise ValueError ( "Only TRAIN and PREDICT modes are supported: %s" % ( mode ) ) return output_spec return model_fn