def _read_tfds(tfds_builder: tfds.core.DatasetBuilder, tfds_split: Text, tfds_skip_decoding_feature: Text, tfds_as_supervised: bool, input_context: Optional[tf.distribute.InputContext] = None, seed: Optional[Union[int, tf.Tensor]] = None, is_training: bool = False, cache: bool = False, cycle_length: Optional[int] = None, block_length: Optional[int] = None) -> tf.data.Dataset: """Reads a dataset from tfds.""" # No op if exist. tfds_builder.download_and_prepare() read_config = tfds.ReadConfig(interleave_cycle_length=cycle_length, interleave_block_length=block_length, input_context=input_context, shuffle_seed=seed) decoders = {} if tfds_skip_decoding_feature: for skip_feature in tfds_skip_decoding_feature.split(','): decoders[skip_feature.strip()] = tfds.decode.SkipDecoding() dataset = tfds_builder.as_dataset(split=tfds_split, shuffle_files=is_training, as_supervised=tfds_as_supervised, decoders=decoders, read_config=read_config) if is_training and not cache: dataset = dataset.repeat() return dataset
def _read_tfds(tfds_builder: tfds.core.DatasetBuilder, tfds_split: Text, tfds_skip_decoding_feature: Text, tfds_as_supervised: bool, input_context: Optional[tf.distribute.InputContext] = None, seed: Optional[Union[int, tf.Tensor]] = None, is_training: bool = False, cache: bool = False, cycle_length: Optional[int] = None, block_length: Optional[int] = None) -> tf.data.Dataset: """Reads a dataset from tfds.""" # No op if exist. tfds_builder.download_and_prepare() decoders = {} if tfds_skip_decoding_feature: for skip_feature in tfds_skip_decoding_feature.split(','): decoders[skip_feature.strip()] = tfds.decode.SkipDecoding() if tfds_builder.info.splits: num_shards = len( tfds_builder.info.splits[tfds_split].file_instructions) else: # The tfds mock path often does not provide splits. num_shards = 1 if input_context and num_shards < input_context.num_input_pipelines: # The number of files in the dataset split is smaller than the number of # input pipelines. We read the entire dataset first and then shard in the # host memory. read_config = tfds.ReadConfig(interleave_cycle_length=cycle_length, interleave_block_length=block_length, input_context=None, shuffle_seed=seed) dataset = tfds_builder.as_dataset(split=tfds_split, shuffle_files=is_training, as_supervised=tfds_as_supervised, decoders=decoders, read_config=read_config) dataset = dataset.shard(input_context.num_input_pipelines, input_context.input_pipeline_id) else: read_config = tfds.ReadConfig(interleave_cycle_length=cycle_length, interleave_block_length=block_length, input_context=input_context, shuffle_seed=seed) dataset = tfds_builder.as_dataset(split=tfds_split, shuffle_files=is_training, as_supervised=tfds_as_supervised, decoders=decoders, read_config=read_config) if is_training and not cache: dataset = dataset.repeat() return dataset
def __init__(self, dataset_builder: tfds.core.DatasetBuilder): """A NitroML dataset from a TFDS DatasetBuilder. Args: dataset_builder: A `tfds.DatasetBuilder` instance which defines the TFDS dataset to use. Example: `dataset = TFDSTask(tfds.builder('titanic'))` """ # TODO(b/159086401): Download and prepare the dataset in a component # instead of at construction time, so that this step happens lazily during # pipeline execution. logging.info('Preparing dataset...') dataset_builder.download_and_prepare() logging.info(dataset_builder.info) self._dataset_builder = dataset_builder self._example_gen = self._make_example_gen()
def _download_and_prepare( args: argparse.Namespace, builder: tfds.core.DatasetBuilder, ) -> None: """Generate a single builder.""" logging.info(f'download_and_prepare for dataset {builder.info.full_name}...') dl_config = _make_download_config(args) if args.add_name_to_manual_dir: dl_config.manual_dir = os.path.join(dl_config.manual_dir, builder.name) builder.download_and_prepare( download_dir=args.download_dir, download_config=dl_config, ) # Dataset generated successfully logging.info('Dataset generation complete...') termcolor.cprint(str(builder.info.as_proto), attrs=['bold'])