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'])