# Copyright 2023 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """A common dataset reader.""" import dataclasses import random from typing import Any, Callable, Dict, List, Optional, Sequence, Text, Union from absl import logging import tensorflow as tf, tf_keras import tensorflow_datasets as tfds from official.core import config_definitions as cfg def _get_random_integer(): return random.randint(0, (1 << 31) - 1) def _maybe_map_fn(dataset: tf.data.Dataset, fn: Optional[Callable[..., Any]] = None) -> tf.data.Dataset: """Calls dataset.map if a valid function is passed in.""" return dataset if fn is None else dataset.map( fn, num_parallel_calls=tf.data.experimental.AUTOTUNE) def match_files(input_path: Union[Sequence[str], str]) -> List[str]: """Matches files from an input_path.""" matched_files = [] # Read dataset from files. usage = ('`input_path` should be either (1) a str indicating a file ' 'path/pattern, or (2) a str indicating multiple file ' 'paths/patterns separated by comma (e.g "a, b, c" or no spaces ' '"a,b,c", or (3) a list of str, each of which is a file ' 'path/pattern or multiple file paths/patterns separated by ' 'comma, but got: %s') if isinstance(input_path, str): input_path_list = [input_path] elif isinstance(input_path, (list, tuple)): if any(not isinstance(x, str) for x in input_path): raise ValueError(usage % input_path) input_path_list = input_path else: raise ValueError(usage % input_path) for input_path in input_path_list: input_patterns = input_path.strip().split(',') for input_pattern in input_patterns: input_pattern = input_pattern.strip() if not input_pattern: continue if '*' in input_pattern or '?' in input_pattern: tmp_matched_files = tf.io.gfile.glob(input_pattern) if not tmp_matched_files: raise ValueError('%s does not match any files.' % input_pattern) matched_files.extend(tmp_matched_files) else: matched_files.append(input_pattern) if not matched_files: raise ValueError('%s does not match any files.' % input_path) return matched_files def _read_files_then_shard(matched_files: List[str], dataset_fn, input_context: Optional[ tf.distribute.InputContext] = None, sharding: bool = False, repeat: bool = False) -> tf.data.Dataset: """Sends all data files to every worker and then shard by data.""" dataset = dataset_fn(matched_files) # When `input_file` is a path to a single file or the number of files is # less than the number of input pipelines, disable auto sharding # so that same input file is sent to all workers. options = tf.data.Options() options.experimental_distribute.auto_shard_policy = ( tf.data.experimental.AutoShardPolicy.OFF) dataset = dataset.with_options(options) # Do not enable sharding if tf.data service is enabled, as sharding will be # handled inside tf.data service. if sharding and input_context and (input_context.num_input_pipelines > 1): dataset = dataset.shard(input_context.num_input_pipelines, input_context.input_pipeline_id) if repeat: dataset = dataset.repeat() return dataset def _shard_files_then_read(matched_files: List[str], dataset_fn, input_context: Optional[ tf.distribute.InputContext] = None, seed: Optional[Union[int, tf.Tensor]] = None, is_training: bool = False, sharding: bool = False, cache: bool = False, cycle_length: Optional[int] = None, block_length: Optional[int] = None, deterministic: bool = False) -> tf.data.Dataset: """Shards the data files and then sent a split to every worker to read.""" dataset = tf.data.Dataset.from_tensor_slices(matched_files) # Shuffle and repeat at file level. # If cache is enabled, `reshuffle_each_iteration` is set to False, # because we will read the same cached data in every iteration anyway. if is_training: # We need a seed to shuffle the files so that when each TPU workers gets # its own shard the files do not overlap. if sharding and seed is None: seed = _get_random_integer() dataset = dataset.shuffle( len(matched_files), seed=seed, reshuffle_each_iteration=True if not cache else False) # Do not enable sharding if tf.data service is enabled, as sharding will be # handled inside tf.data service. if sharding and input_context and (input_context.num_input_pipelines > 1): dataset = dataset.shard(input_context.num_input_pipelines, input_context.input_pipeline_id) # If cache is enabled, we will call `repeat()` later after `cache()`. if is_training and not cache: dataset = dataset.repeat() dataset = dataset.interleave( map_func=dataset_fn, cycle_length=cycle_length, block_length=block_length, num_parallel_calls=(cycle_length if cycle_length else tf.data.experimental.AUTOTUNE), deterministic=deterministic) return dataset def _read_tfds(tfds_name: Text, tfds_data_dir: Text, 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.""" repeat_filenames = is_training and not cache read_config = tfds.ReadConfig( interleave_cycle_length=cycle_length, interleave_block_length=block_length, input_context=input_context, shuffle_seed=seed, repeat_filenames=repeat_filenames, # Only assert cardinality when we have a finite dataset. assert_cardinality=not repeat_filenames, skip_prefetch=True) decoders = {} if tfds_skip_decoding_feature: for skip_feature in tfds_skip_decoding_feature.split(','): decoders[skip_feature.strip()] = tfds.decode.SkipDecoding() if tfds_name.startswith('mldataset.'): dataset = tfds.load(name=tfds_name, split=tfds_split, as_supervised=tfds_as_supervised, decoders=decoders if decoders else None, read_config=read_config) else: builder = tfds.builder(tfds_name, data_dir=tfds_data_dir) if builder.info.splits: num_shards = len(builder.info.splits[tfds_split].file_instructions) else: # The tfds mock path often does not provide splits. num_shards = 1 load_kwargs = dict( name=tfds_name, download=True, split=tfds_split, shuffle_files=is_training, as_supervised=tfds_as_supervised, decoders=decoders if decoders else None) if tfds_data_dir: load_kwargs.update({'data_dir': tfds_data_dir}) 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 = dataclasses.replace(read_config, input_context=None) load_kwargs.update({'read_config': read_config}) dataset = tfds.load(**load_kwargs) dataset = dataset.shard(input_context.num_input_pipelines, input_context.input_pipeline_id) else: load_kwargs.update({'read_config': read_config}) dataset = tfds.load(**load_kwargs) return dataset class InputReader: """Input reader that returns a tf.data.Dataset instance.""" # A static random number which is the same across different InputReader # instances. static_randnum = _get_random_integer() def __init__( self, params: cfg.DataConfig, dataset_fn=tf.data.TFRecordDataset, decoder_fn: Optional[Callable[..., Any]] = None, combine_fn: Optional[Callable[..., Any]] = None, sample_fn: Optional[Callable[..., Any]] = None, parser_fn: Optional[Callable[..., Any]] = None, filter_fn: Optional[Callable[..., tf.Tensor]] = None, transform_and_batch_fn: Optional[ Callable[ [tf.data.Dataset, Optional[tf.distribute.InputContext]], tf.data.Dataset, ] ] = None, postprocess_fn: Optional[Callable[..., Any]] = None, ): """Initializes an InputReader instance. Args: params: A config_definitions.DataConfig object. dataset_fn: A `tf.data.Dataset` that consumes the input files. For example, it can be `tf.data.TFRecordDataset`. decoder_fn: An optional `callable` that takes the serialized data string and decodes them into the raw tensor dictionary. combine_fn: An optional `callable` that takes a dictionarty of `tf.data.Dataset` objects as input and outputs a combined dataset. It will be executed after the decoder_fn and before the sample_fn. sample_fn: An optional `callable` that takes a `tf.data.Dataset` object as input and outputs the transformed dataset. It performs sampling on the decoded raw tensors dict before the parser_fn. parser_fn: An optional `callable` that takes the decoded raw tensors dict and parse them into a dictionary of tensors that can be consumed by the model. It will be executed after decoder_fn. filter_fn: An optional `callable` mapping a dataset element to a boolean. It will be executed after parser_fn. transform_and_batch_fn: An optional `callable` that takes a `tf.data.Dataset` object and an optional `tf.distribute.InputContext` as input, and returns a `tf.data.Dataset` object. It will be executed after `parser_fn` to transform and batch the dataset; if None, after `parser_fn` is executed, the dataset will be batched into per-replica batch size. postprocess_fn: A optional `callable` that processes batched tensors. It will be executed after batching. """ if params.input_path and params.tfds_name: raise ValueError('At most one of `input_path` and `tfds_name` can be ' 'specified, but got %s and %s.' % (params.input_path, params.tfds_name)) if (isinstance(params.input_path, cfg.base_config.Config) or isinstance(params.tfds_name, cfg.base_config.Config) ) and combine_fn is None: raise ValueError( 'A combine_fn is required if `input_path` or `tfds_name` is a dict.') self._tfds_name = params.tfds_name self._tfds_data_dir = params.tfds_data_dir self._matched_files = None if not params.input_path: # Read dataset from TFDS. if not params.tfds_split: raise ValueError( '`tfds_name` is %s, but `tfds_split` is not specified.' % params.tfds_name) else: self._matched_files = self.get_files(params.input_path) self._global_batch_size = params.global_batch_size self._is_training = params.is_training self._drop_remainder = params.drop_remainder self._shuffle_buffer_size = params.shuffle_buffer_size self._cache = params.cache self._cycle_length = params.cycle_length self._block_length = params.block_length self._deterministic = params.deterministic self._sharding = params.sharding self._tfds_split = params.tfds_split self._tfds_as_supervised = params.tfds_as_supervised self._tfds_skip_decoding_feature = params.tfds_skip_decoding_feature self._dataset_fn = dataset_fn self._decoder_fn = decoder_fn self._combine_fn = combine_fn self._sample_fn = sample_fn self._parser_fn = parser_fn self._transform_and_batch_fn = transform_and_batch_fn self._postprocess_fn = postprocess_fn self._filter_fn = filter_fn self._seed = params.seed self._prefetch_buffer_size = ( params.prefetch_buffer_size or tf.data.experimental.AUTOTUNE) self._autotune_algorithm = params.autotune_algorithm # When tf.data service is enabled, each data service worker should get # different random seeds. Thus, we set `seed` to None. # Sharding should also be disabled because tf data service handles how # each worker shard data with `processing_mode` in distribute method. if params.enable_tf_data_service: self._seed = None self._sharding = False self._enable_tf_data_service = ( params.enable_tf_data_service and params.tf_data_service_address) self._tf_data_service_address = params.tf_data_service_address self._enable_shared_tf_data_service_between_parallel_trainers = ( params.enable_shared_tf_data_service_between_parallel_trainers) self._apply_tf_data_service_before_batching = ( params.apply_tf_data_service_before_batching) self._trainer_id = params.trainer_id if self._enable_tf_data_service: # Add a random seed as the tf.data service job name suffix, so tf.data # service doesn't reuse the previous state if TPU worker gets preempted. # It's necessary to add global batch size into the tf data service job # name because when tuning batch size with vizier and tf data service is # also enable, the tf data servce job name should be different for # different vizier trials since once batch size is changed, from the # tf.data perspective, the dataset is a different instance, and a # different job name should be used for tf data service. Otherwise, the # model would read tensors from the incorrect tf data service job, which # would causes dimension mismatch on the batch size dimension. self._tf_data_service_job_name = ( f'{params.tf_data_service_job_name}_bs{params.global_batch_size}_' f'{self.static_randnum}') self._enable_round_robin_tf_data_service = params.get( 'enable_round_robin_tf_data_service', False) if self._enable_shared_tf_data_service_between_parallel_trainers: # When shared tf.data service is enabled, only a single tf.data service # instance should be created and shared between parallel trainers. If # the global batch size is different across trainers, # params.apply_tf_data_service_before_batching should be set to true # because tf.data service with different batch sizes will be considered # separate tf.data service instances. self._tf_data_service_job_name = ( f'{params.tf_data_service_job_name}_{self.static_randnum}') def get_files(self, input_path): """Gets matched files. Can be overridden by subclasses.""" if not input_path: return None # we want to combine / mix datasets if isinstance(input_path, cfg.base_config.Config): matched_files = {} for k, v in input_path.as_dict().items(): matched_files[k] = match_files(v) # single dataset else: matched_files = match_files(input_path) return matched_files def _read_data_source( self, matched_files: Union[Dict[str, List[str]], List[str]], dataset_fn, input_context: Optional[tf.distribute.InputContext] = None, ): """Reads the data source (files/tfds) to a dataset.""" def _files_to_dataset(files: List[str]) -> tf.data.Dataset: if len(files) > 1: if input_context and (len(files) < input_context.num_input_pipelines): logging.warn( ( 'The number of files %d is less than the number of input ' 'pipelines %d. We will send all input files to every worker. ' 'Please consider sharding your data into more files.' ), len(files), input_context.num_input_pipelines, ) return _read_files_then_shard( files, dataset_fn, input_context, sharding=self._sharding, repeat=self._is_training and not self._cache) else: return _shard_files_then_read( files, dataset_fn, input_context, seed=self._seed, is_training=self._is_training, sharding=self._sharding, cache=self._cache, cycle_length=self._cycle_length, block_length=self._block_length, deterministic=self._deterministic) elif len(files) == 1: return _read_files_then_shard( files, dataset_fn, input_context, sharding=self._sharding, repeat=self._is_training and not self._cache) else: raise ValueError('It is unexpected that `tfds_builder` is None and ' 'there is also no `files`.') if self._tfds_name: if isinstance(self._tfds_name, cfg.base_config.Config): dataset = {} for k, tfds_name in self._tfds_name.as_dict().items(): dataset[k] = _read_tfds( tfds_name=tfds_name, tfds_data_dir=self._tfds_data_dir, tfds_split=self._tfds_split, tfds_skip_decoding_feature=self._tfds_skip_decoding_feature, tfds_as_supervised=self._tfds_as_supervised, input_context=input_context, seed=self._seed, is_training=self._is_training, cache=self._cache, cycle_length=self._cycle_length, block_length=self._block_length) else: dataset = _read_tfds( tfds_name=self._tfds_name, tfds_data_dir=self._tfds_data_dir, tfds_split=self._tfds_split, tfds_skip_decoding_feature=self._tfds_skip_decoding_feature, tfds_as_supervised=self._tfds_as_supervised, input_context=input_context, seed=self._seed, is_training=self._is_training, cache=self._cache, cycle_length=self._cycle_length, block_length=self._block_length) elif isinstance(matched_files, (list, tuple)): dataset = _files_to_dataset(matched_files) elif isinstance(matched_files, dict): dataset = {} for k, fs in matched_files.items(): dataset[k] = _files_to_dataset(fs) else: raise ValueError('`matched_files` should be a list or dict.') return dataset def _decode_and_parse_dataset( self, dataset: Union[tf.data.Dataset, Dict[Text, tf.data.Dataset]], batch_size: int, input_context: Optional[tf.distribute.InputContext] = None ) -> tf.data.Dataset: """Returns a tf.data.Dataset object after shuffling, decoding, and parsing.""" def _shuffle_and_decode(ds): # If cache is enabled, we will call `shuffle()` later after `cache()`. if self._is_training and not self._cache: ds = ds.shuffle(self._shuffle_buffer_size, seed=self._seed) # Decode ds = _maybe_map_fn(ds, self._decoder_fn) return ds dataset = tf.nest.map_structure(_shuffle_and_decode, dataset) if tf.nest.is_nested(dataset): dataset = self._combine_fn(dataset) if self._sample_fn is not None: dataset = dataset.apply(self._sample_fn) dataset = _maybe_map_fn(dataset, self._parser_fn) if self._filter_fn is not None: dataset = dataset.filter(self._filter_fn) if self._cache: dataset = dataset.cache() if self._is_training: dataset = dataset.repeat() dataset = dataset.shuffle(self._shuffle_buffer_size, seed=self._seed) # Applies tf.data service before batching operations. This is useful when # tf.data service is shared between parallel trainers, and batch size is # changing between parallel trainers. Then batch size is changing, tf.data # services will be considered different instances if applied after batching # operations, which make it difficult to share between parallel trainers. # However, if there are additional expensive operations in # self._transform_and_batch_fn and self._postprocess_fn, the entire tf.data # pipeline could be slowed down. In this case, try to move these dataset # operations into early stages if possible. if (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._transform_and_batch_fn is not None: dataset = self._transform_and_batch_fn(dataset, input_context) else: per_replica_batch_size = input_context.get_per_replica_batch_size( batch_size) if input_context else batch_size dataset = dataset.batch( per_replica_batch_size, drop_remainder=self._drop_remainder) return dataset def _maybe_apply_data_service( self, dataset: tf.data.Dataset, input_context: Optional[tf.distribute.InputContext] = None ) -> tf.data.Dataset: """Potentially distributes a dataset.""" if self._enable_tf_data_service and input_context: if self._enable_round_robin_tf_data_service: replicas_per_input_pipeline = input_context.num_replicas_in_sync // ( input_context.num_input_pipelines) base_consumer_index = input_context.input_pipeline_id * ( replicas_per_input_pipeline) num_consumers = input_context.num_input_pipelines * ( replicas_per_input_pipeline) range_dataset = tf.data.Dataset.range(replicas_per_input_pipeline) tfds_kwargs = { 'processing_mode': 'parallel_epochs', 'service': self._tf_data_service_address, 'job_name': self._tf_data_service_job_name, 'num_consumers': num_consumers } if self._enable_shared_tf_data_service_between_parallel_trainers: raise ValueError('Shared tf.data service does not support round-robin' ' tf.data service.') dataset = range_dataset.map(lambda i: dataset.apply( # pylint: disable=g-long-lambda tf.data.experimental.service.distribute( consumer_index=base_consumer_index + i, **tfds_kwargs))) # Use parallel interleave to read multiple batches from a tf.data # service worker in parallel. dataset = dataset.interleave( lambda x: x, cycle_length=replicas_per_input_pipeline, num_parallel_calls=replicas_per_input_pipeline, deterministic=True) else: tfds_kwargs = { 'processing_mode': 'parallel_epochs', 'service': self._tf_data_service_address, 'job_name': self._tf_data_service_job_name, } if self._enable_shared_tf_data_service_between_parallel_trainers: tfds_kwargs.update({ 'processing_mode': tf.data.experimental.service.ShardingPolicy.OFF, 'cross_trainer_cache': tf.data.experimental.service.CrossTrainerCache( trainer_id=self._trainer_id) }) dataset = dataset.apply( tf.data.experimental.service.distribute(**tfds_kwargs)) 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) return dataset.prefetch(self._prefetch_buffer_size)