(bert_config, num_labels, 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, num_labels, 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))

    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]
    label_ids = features["label_ids"]
    is_real_example = None
    if "is_real_example" in features:
      is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32)
    else:
      is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)

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

    (total_loss, per_example_loss, logits, probabilities) = create_model(
        bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
        num_labels, 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:

      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.EVAL:

      def metric_fn(per_example_loss, label_ids, logits, is_real_example):
        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
        accuracy = tf.metrics.accuracy(
            labels=label_ids, predictions=predictions, weights=is_real_example)
        loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example)
        return {
            "eval_accuracy": accuracy,
            "eval_loss": loss,
        }

      eval_metrics = (metric_fn,
                      [per_example_loss, label_ids, logits, is_real_example])
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          eval_metrics=eval_metrics,
          scaffold_fn=scaffold_fn)
    else:
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          predictions={"probabilities": probabilities},
          scaffold_fn=scaffold_fn)
    return output_spec

  return model_fn
def model_fn_builder ( bert_config , num_labels , 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 ) ) input_ids = features [ "input_ids" ] input_mask = features [ "input_mask" ] segment_ids = features [ "segment_ids" ] label_ids = features [ "label_ids" ] is_real_example = None if "is_real_example" in features : is_real_example = tf . cast ( features [ "is_real_example" ] , dtype = tf . float32 ) else : is_real_example = tf . ones ( tf . shape ( label_ids ) , dtype = tf . float32 ) is_training = ( mode == tf . estimator . ModeKeys . TRAIN ) ( total_loss , per_example_loss , logits , probabilities ) = create_model ( bert_config , is_training , input_ids , input_mask , segment_ids , label_ids , num_labels , 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 : 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 . EVAL : def metric_fn ( per_example_loss , label_ids , logits , is_real_example ) : predictions = tf . argmax ( logits , axis = - 1 , output_type = tf . int32 ) accuracy = tf . metrics . accuracy ( labels = label_ids , predictions = predictions , weights = is_real_example ) loss = tf . metrics . mean ( values = per_example_loss , weights = is_real_example ) return { "eval_accuracy" : accuracy , "eval_loss" : loss , } eval_metrics = ( metric_fn , [ per_example_loss , label_ids , logits , is_real_example ] ) output_spec = tf . contrib . tpu . TPUEstimatorSpec ( mode = mode , loss = total_loss , eval_metrics = eval_metrics , scaffold_fn = scaffold_fn ) else : output_spec = tf . contrib . tpu . TPUEstimatorSpec ( mode = mode , predictions = { "probabilities" : probabilities } , scaffold_fn = scaffold_fn ) return output_spec return model_fn