(input_glob, seq_length, is_training, drop_remainder,
                     num_hosts, num_threads=8)
return input_fn
Creates an `input_fn` closure to be passed to TPUEstimator.
Creates an `input_fn` closure to be passed to TPUEstimator.
Creates an input_fn closure to be passed to TPUEstimator .
def input_fn_builder(input_glob, seq_length, is_training, drop_remainder,
                     num_hosts, num_threads=8):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.float32),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "cls_index": tf.FixedLenFeature([], tf.int64),
      "p_mask": tf.FixedLenFeature([seq_length], tf.float32)
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["is_impossible"] = tf.FixedLenFeature([], tf.float32)

  tf.logging.info("Input tfrecord file glob {}".format(input_glob))
  global_input_paths = tf.gfile.Glob(input_glob)
  tf.logging.info("Find {} input paths {}".format(
      len(global_input_paths), global_input_paths))

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    if FLAGS.use_tpu:
      batch_size = params["batch_size"]
    elif is_training:
      batch_size = FLAGS.train_batch_size
    else:
      batch_size = FLAGS.predict_batch_size

    # Split tfrecords across hosts
    if num_hosts > 1:
      host_id = params["context"].current_host
      num_files = len(global_input_paths)
      if num_files >= num_hosts:
        num_files_per_host = (num_files + num_hosts - 1) // num_hosts
        my_start_file_id = host_id * num_files_per_host
        my_end_file_id = min((host_id + 1) * num_files_per_host, num_files)
        input_paths = global_input_paths[my_start_file_id: my_end_file_id]
      tf.logging.info("Host {} handles {} files".format(host_id,
                                                        len(input_paths)))
    else:
      input_paths = global_input_paths

    if len(input_paths) == 1:
      d = tf.data.TFRecordDataset(input_paths[0])
      # For training, we want a lot of parallel reading and shuffling.
      # For eval, we want no shuffling and parallel reading doesn't matter.
      if is_training:
        d = d.shuffle(buffer_size=FLAGS.shuffle_buffer)
        d = d.repeat()
    else:
      d = tf.data.Dataset.from_tensor_slices(input_paths)
      # file level shuffle
      d = d.shuffle(len(input_paths)).repeat()

      # `cycle_length` is the number of parallel files that get read.
      cycle_length = min(num_threads, len(input_paths))

      d = d.apply(
          tf.contrib.data.parallel_interleave(
              tf.data.TFRecordDataset,
              sloppy=is_training,
              cycle_length=cycle_length))

      if is_training:
        # sample level shuffle
        d = d.shuffle(buffer_size=FLAGS.shuffle_buffer)

    d = d.apply(
        tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            num_parallel_batches=num_threads,
            drop_remainder=drop_remainder))
    d = d.prefetch(1024)

    return d

  return input_fn
def input_fn_builder ( input_glob , seq_length , is_training , drop_remainder , num_hosts , num_threads = 8 ) : name_to_features = { "unique_ids" : tf . FixedLenFeature ( [ ] , tf . int64 ) , "input_ids" : tf . FixedLenFeature ( [ seq_length ] , tf . int64 ) , "input_mask" : tf . FixedLenFeature ( [ seq_length ] , tf . float32 ) , "segment_ids" : tf . FixedLenFeature ( [ seq_length ] , tf . int64 ) , "cls_index" : tf . FixedLenFeature ( [ ] , tf . int64 ) , "p_mask" : tf . FixedLenFeature ( [ seq_length ] , tf . float32 ) } if is_training : name_to_features [ "start_positions" ] = tf . FixedLenFeature ( [ ] , tf . int64 ) name_to_features [ "end_positions" ] = tf . FixedLenFeature ( [ ] , tf . int64 ) name_to_features [ "is_impossible" ] = tf . FixedLenFeature ( [ ] , tf . float32 ) tf . logging . info ( "Input tfrecord file glob {}" . format ( input_glob ) ) global_input_paths = tf . gfile . Glob ( input_glob ) tf . logging . info ( "Find {} input paths {}" . format ( len ( global_input_paths ) , global_input_paths ) ) def _decode_record ( record , name_to_features ) : """Decodes a record to a TensorFlow example.""" example = tf . parse_single_example ( record , name_to_features ) # tf.Example only supports tf.int64, but the TPU only supports tf.int32. # So cast all int64 to int32. for name in list ( example . keys ( ) ) : t = example [ name ] if t . dtype == tf . int64 : t = tf . cast ( t , tf . int32 ) example [ name ] = t return example def input_fn ( params ) : """The actual input function.""" if FLAGS . use_tpu : batch_size = params [ "batch_size" ] elif is_training : batch_size = FLAGS . train_batch_size else : batch_size = FLAGS . predict_batch_size # Split tfrecords across hosts if num_hosts > 1 : host_id = params [ "context" ] . current_host num_files = len ( global_input_paths ) if num_files >= num_hosts : num_files_per_host = ( num_files + num_hosts - 1 ) // num_hosts my_start_file_id = host_id * num_files_per_host my_end_file_id = min ( ( host_id + 1 ) * num_files_per_host , num_files ) input_paths = global_input_paths [ my_start_file_id : my_end_file_id ] tf . logging . info ( "Host {} handles {} files" . format ( host_id , len ( input_paths ) ) ) else : input_paths = global_input_paths if len ( input_paths ) == 1 : d = tf . data . TFRecordDataset ( input_paths [ 0 ] ) # For training, we want a lot of parallel reading and shuffling. # For eval, we want no shuffling and parallel reading doesn't matter. if is_training : d = d . shuffle ( buffer_size = FLAGS . shuffle_buffer ) d = d . repeat ( ) else : d = tf . data . Dataset . from_tensor_slices ( input_paths ) # file level shuffle d = d . shuffle ( len ( input_paths ) ) . repeat ( ) # `cycle_length` is the number of parallel files that get read. cycle_length = min ( num_threads , len ( input_paths ) ) d = d . apply ( tf . contrib . data . parallel_interleave ( tf . data . TFRecordDataset , sloppy = is_training , cycle_length = cycle_length ) ) if is_training : # sample level shuffle d = d . shuffle ( buffer_size = FLAGS . shuffle_buffer ) d = d . apply ( tf . contrib . data . map_and_batch ( lambda record : _decode_record ( record , name_to_features ) , batch_size = batch_size , num_parallel_batches = num_threads , drop_remainder = drop_remainder ) ) d = d . prefetch ( 1024 ) return d return input_fn