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 create_split(dataset_builder: tfds.core.DatasetBuilder, batch_size: int, train: bool, dtype: tf.DType = tf.float32, image_size: int = IMAGE_SIZE, cache: bool = False): """Creates a split from the ImageNet dataset using TensorFlow Datasets. Args: dataset_builder: TFDS dataset builder for ImageNet. batch_size: the batch size returned by the data pipeline. train: Whether to load the train or evaluation split. dtype: data type of the image (default: float32). image_size: The target size of the images (default: 224). cache: Whether to cache the dataset (default: False). Returns: A `tf.data.Dataset`. """ if train: train_size = dataset_builder.info.splits['train'].num_examples split_size = train_size // jax.host_count() start = jax.host_id() * split_size split = 'train[{}:{}]'.format(start, start + split_size) else: validation_size = dataset_builder.info.splits[ 'validation'].num_examples split_size = validation_size // jax.host_count() start = jax.host_id() * split_size split = 'validation[{}:{}]'.format(start, start + split_size) def _decode_example(example): if train: image = preprocess_for_train(example['image'], dtype, image_size) else: image = preprocess_for_eval(example['image'], dtype, image_size) return {'image': image, 'label': example['label']} ds = dataset_builder.as_dataset( split=split, decoders={'image': tfds.decode.SkipDecoding()}) ds.options().experimental_threading.private_threadpool_size = 48 ds.options().experimental_threading.max_intra_op_parallelism = 1 if cache: ds = ds.cache() if train: ds = ds.repeat() ds = ds.shuffle(16 * batch_size, seed=0) ds = ds.map(_decode_example, num_parallel_calls=tf.data.experimental.AUTOTUNE) ds = ds.batch(batch_size, drop_remainder=True) if not train: ds = ds.repeat() ds = ds.prefetch(10) return ds
def get_raw_dataset(dataset_builder: tfds.core.DatasetBuilder, split: str, *, reverse_translation: bool = False) -> tf.data.Dataset: """Loads a raw WMT dataset and normalizes feature keys. Args: dataset_builder: TFDS dataset builder that can build `slit`. split: Split to use. This must be the full split. We shard the split across multiple hosts and currently don't support sharding subsplits. reverse_translation: bool: whether to reverse the translation direction. e.g. for 'de-en' this translates from english to german. Returns: Dataset with source and target language features mapped to 'inputs' and 'targets'. """ num_examples = dataset_builder.info.splits[split].num_examples per_host_split = deterministic_data.get_read_instruction_for_host( split, num_examples, drop_remainder=False) ds = dataset_builder.as_dataset(split=per_host_split, shuffle_files=False) ds = ds.map(NormalizeFetaureNamesOp( dataset_builder.info, reverse_translation=reverse_translation), num_parallel_calls=AUTOTUNE) return ds
def _build_autocached_info(self, builder: tfds.core.DatasetBuilder): """Returns the auto-cache information string.""" always_cached = {} never_cached = {} unshuffle_cached = {} for split_name in sorted(builder.info.splits.keys()): split_name = str(split_name) cache_shuffled = builder._should_cache_ds( # pylint: disable=protected-access split_name, shuffle_files=True, read_config=tfds.ReadConfig()) cache_unshuffled = builder._should_cache_ds( # pylint: disable=protected-access split_name, shuffle_files=False, read_config=tfds.ReadConfig()) if all((cache_shuffled, cache_unshuffled)): always_cached[split_name] = None elif not any((cache_shuffled, cache_unshuffled)): never_cached[split_name] = None else: # Dataset is only cached when shuffled_files is False assert not cache_shuffled and cache_unshuffled unshuffle_cached[split_name] = None if not len(builder.info.splits) or not builder.info.dataset_size: # pylint: disable=g-explicit-length-test autocached_info = 'Unknown' elif len(always_cached) == len(builder.info.splits.keys()): autocached_info = 'Yes' # All splits are auto-cached. elif len(never_cached) == len(builder.info.splits.keys()): autocached_info = 'No' # Splits never auto-cached. else: # Some splits cached, some not. autocached_info_parts = [] if always_cached: split_names_str = ', '.join(always_cached) autocached_info_parts.append( 'Yes ({})'.format(split_names_str)) if never_cached: split_names_str = ', '.join(never_cached) autocached_info_parts.append('No ({})'.format(split_names_str)) if unshuffle_cached: split_names_str = ', '.join(unshuffle_cached) autocached_info_parts.append( 'Only when `shuffle_files=False` ({})'.format( split_names_str)) autocached_info = ', '.join(autocached_info_parts) return autocached_info
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'])
def get_raw_dataset(dataset_builder: tfds.core.DatasetBuilder, split: str) -> tf.data.Dataset: """Loads a raw text dataset and normalizes feature keys. Args: dataset_builder: TFDS dataset builder that can build `split`. split: Split to use. This must be the full split. We shard the split across multiple hosts and currently don't support sharding subsplits. Returns: Dataset with source and target language features mapped to 'inputs' and 'targets'. """ per_host_split = deterministic_data.get_read_instruction_for_host( split, dataset_info=dataset_builder.info, drop_remainder=False) ds = dataset_builder.as_dataset(split=per_host_split, shuffle_files=False) ds = ds.map(NormalizeFeatureNamesOp(dataset_builder.info), num_parallel_calls=AUTOTUNE) return ds
def create_split( dataset_builder: tfds.core.DatasetBuilder, batch_size: int, train: bool = True, half_precision: bool = False, image_size: int = IMAGE_SIZE, mean: Optional[Tuple[float]] = None, std: Optional[Tuple[float]] = None, interpolation: str = 'bicubic', augment_name: Optional[str] = None, randaug_num_layers: Optional[int] = None, randaug_magnitude: Optional[int] = None, cache: bool = False, no_repeat: bool = False, ): """Creates a split from the ImageNet dataset using TensorFlow Datasets. Args: dataset_builder: TFDS dataset builder for ImageNet. batch_size: the batch size returned by the data pipeline. train: Whether to load the train or evaluation split. half_precision: convert image datatype to half-precision image_size: The target size of the images (default: 224). mean: image dataset mean std: image dataset std-dev interpolation: interpolation method to use for image resize (default: 'bicubic') cache: Whether to cache the dataset (default: False). no_repeat: disable repeat iter for evaluation Returns: A `tf.data.Dataset`. """ mean = mean or MEAN_RGB std = std or STDDEV_RGB interpolation = tf.image.ResizeMethod.BICUBIC if interpolation == 'bicubic' else tf.image.ResizeMethod.BILINEAR platform = jax.local_devices()[0].platform if half_precision: if platform == 'tpu': input_dtype = tf.bfloat16 else: input_dtype = tf.float16 else: input_dtype = tf.float32 if train: data_size = dataset_builder.info.splits['train'].num_examples split = 'train' else: data_size = dataset_builder.info.splits['validation'].num_examples split = 'validation' split_size = data_size // jax.host_count() start = jax.host_id() * split_size split = split + '[{}:{}]'.format(start, start + split_size) def _decode_example(example): if train: image = preprocess_for_train(example['image'], input_dtype, image_size, mean, std, interpolation, augment_name=augment_name, randaug_num_layers=randaug_num_layers, randaug_magnitude=randaug_magnitude) else: image = preprocess_for_eval(example['image'], input_dtype, image_size, mean, std, interpolation) return {'image': image, 'label': example['label']} ds = dataset_builder.as_dataset( split=split, decoders={'image': tfds.decode.SkipDecoding()}) ds.options().experimental_threading.private_threadpool_size = 16 ds.options().experimental_threading.max_intra_op_parallelism = 1 if cache: ds = ds.cache() if train: ds = ds.repeat() ds = ds.shuffle(16 * batch_size, seed=0) ds = ds.map(_decode_example, num_parallel_calls=tf.data.experimental.AUTOTUNE) ds = ds.batch(batch_size, drop_remainder=True) if not train and not no_repeat: ds = ds.repeat() ds = ds.prefetch(10) return ds