Method read

official/core/input_reader.py:569–599

Generates a tf.data.Dataset object.

(self,
           input_context: Optional[tf.distribute.InputContext] = None,
           dataset: Optional[tf.data.Dataset] = None)

    return dataset

  def read(self,
           input_context: Optional[tf.distribute.InputContext] = None,
           dataset: Optional[tf.data.Dataset] = None) -> tf.data.Dataset:
    """Generates a tf.data.Dataset object."""
    if dataset is None:
      dataset = self._read_data_source(self._matched_files, self._dataset_fn,
                                       input_context)
    dataset = self._decode_and_parse_dataset(dataset, self._global_batch_size,
                                             input_context)
    dataset = _maybe_map_fn(dataset, self._postprocess_fn)
    if not (self._enable_shared_tf_data_service_between_parallel_trainers and
            self._apply_tf_data_service_before_batching):
      dataset = self._maybe_apply_data_service(dataset, input_context)

    if self._deterministic is not None:
      options = tf.data.Options()
      options.deterministic = self._deterministic
      dataset = dataset.with_options(options)
    if self._autotune_algorithm:
      options = tf.data.Options()
      options.autotune.autotune_algorithm = (
          tf.data.experimental.AutotuneAlgorithm[self._autotune_algorithm]
      )
      dataset = dataset.with_options(options)

    if self._ram_budget:
      options = tf.data.Options()
      options.autotune.ram_budget = self._ram_budget * 1024 * 1024 * 1024
      dataset = dataset.with_options(options)

    return dataset.prefetch(self._prefetch_buffer_size)