|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Dataset reader for vision model garden."""
|
|
|
| from typing import Any, Callable, Mapping, Optional, Tuple, Union
|
|
|
| from absl import logging
|
| import tensorflow as tf, tf_keras
|
|
|
| from official.core import config_definitions as cfg
|
| from official.core import input_reader
|
|
|
|
|
| InputReader = input_reader.InputReader
|
|
|
|
|
| def build_weighted_sampling_combine_fn(
|
| weights: Mapping[Any, Any], stop_on_empty_dataset=True
|
| ) -> Callable[[tf.data.Dataset], tf.data.Dataset]:
|
| """Builds a combine_fn using weighted sampling."""
|
|
|
| def combine_fn(datasets: Mapping[Any, tf.data.Dataset]) -> tf.data.Dataset:
|
| """Combines multiple datasets using weighted sampling."""
|
| ds = []
|
| ws = []
|
| for k, dataset in datasets.items():
|
| ds.append(dataset)
|
| ws.append(weights[k])
|
| return tf.data.Dataset.sample_from_datasets(
|
| ds, ws, stop_on_empty_dataset=stop_on_empty_dataset)
|
|
|
| return combine_fn
|
|
|
|
|
| def create_combine_fn(
|
| params: cfg.DataConfig
|
| ) -> Union[None, Callable[[tf.data.Dataset], tf.data.Dataset]]:
|
| """Creates and returns a combine_fn for dataset mixing."""
|
| if (
|
| hasattr(params, 'stop_on_empty_dataset')
|
| and params.stop_on_empty_dataset is not None
|
| ):
|
| stop_on_empty_dataset = params.stop_on_empty_dataset
|
| else:
|
| stop_on_empty_dataset = True
|
|
|
| if params.weights:
|
|
|
| if (not isinstance(params.input_path, cfg.base_config.Config) or
|
| not isinstance(params.weights, cfg.base_config.Config)):
|
| raise ValueError(
|
| 'input_path and weights must both be a Config to use weighted '
|
| 'sampling.')
|
| input_paths = params.input_path.as_dict()
|
| weights = params.weights.as_dict()
|
| if len(input_paths) != len(weights):
|
| raise ValueError(
|
| 'The number of input_path and weights must be the same, but got %d '
|
| 'input_paths and %d weights.' % (len(input_paths), len(weights)))
|
|
|
| for k in input_paths.keys():
|
| if k not in weights:
|
| raise ValueError(
|
| 'input_path key \'%s\' does not have a corresponding weight.' % k)
|
|
|
| return build_weighted_sampling_combine_fn(weights, stop_on_empty_dataset)
|
| return None
|
|
|
|
|
| def calculate_batch_sizes(total_batch_size: int,
|
| pseudo_label_ratio: float,
|
| pseudo_label_batch_size: int = 0) -> Tuple[int, int]:
|
| """Calculates labeled and pseudo-labeled dataset batch sizes.
|
|
|
| Returns (labeled_batch_size, pseudo_labeled_batch_size) given a
|
| total batch size and pseudo-label data ratio.
|
|
|
| Args:
|
| total_batch_size: The total batch size for all data.
|
| pseudo_label_ratio: A float ratio of pseudo-labeled to labeled data in a
|
| batch. If it is negative, use `pseudo_label_batch_size` instead.
|
| pseudo_label_batch_size: The batch size of pseudo-labeled data. It is ignored
|
| if `pseudo_label_ratio` is valid. If not, it will be used and it cannot be
|
| larger than total global batch size or less than 0 if pseudo_label_ratio is
|
| also less than 0.
|
|
|
| Returns:
|
| (labeled_batch_size, pseudo_labeled_batch_size) as ints.
|
|
|
| Raises:
|
| ValueError: If total_batch_size is negative, or both If pseudo_label_ratio
|
| is negative and pseudo-label global_batch_size is negative or larger than
|
| total batch size.
|
| """
|
| if total_batch_size < 0:
|
| raise ValueError('Invalid total_batch_size: {}'.format(total_batch_size))
|
| if pseudo_label_ratio >= 0.0:
|
| ratio_factor = pseudo_label_ratio / (1.0 + pseudo_label_ratio)
|
| pseudo_label_batch_size = int(total_batch_size * ratio_factor)
|
| label_batch_size = total_batch_size - pseudo_label_batch_size
|
| else:
|
| if pseudo_label_batch_size > total_batch_size or pseudo_label_batch_size < 0:
|
| raise ValueError(
|
| 'The batch size of pseudo-label dataset should not be larger than '
|
| 'total global batch size.')
|
| logging.info('data_ratio for pseudo-label dataset is less than 0. '
|
| 'Use global_batch_size from pseudo_label data config instead.')
|
| label_batch_size = total_batch_size - pseudo_label_batch_size
|
| return label_batch_size, pseudo_label_batch_size
|
|
|
|
|
| class CombinationDatasetInputReader(input_reader.InputReader):
|
| """Combination dataset input reader."""
|
|
|
| def __init__(self,
|
| params: cfg.DataConfig,
|
| dataset_fn=tf.data.TFRecordDataset,
|
| pseudo_label_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,
|
| 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 CombinationDatasetInputReader instance.
|
|
|
| This class mixes a labeled and pseudo-labeled dataset. The params
|
| must contain "pseudo_label_data.input_path" to specify the
|
| pseudo-label dataset files and "pseudo_label_data.data_ratio"
|
| to specify a per-batch mixing ratio of pseudo-label examples to
|
| labeled dataset examples.
|
|
|
| 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`.
|
| pseudo_label_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.
|
| 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.
|
|
|
| Raises:
|
| ValueError: If drop_remainder is False.
|
| """
|
| super().__init__(
|
| params=params,
|
| dataset_fn=dataset_fn,
|
| decoder_fn=decoder_fn,
|
| combine_fn=combine_fn,
|
| sample_fn=sample_fn,
|
| parser_fn=parser_fn,
|
| transform_and_batch_fn=transform_and_batch_fn,
|
| postprocess_fn=postprocess_fn)
|
|
|
| self._pseudo_label_file_pattern = params.pseudo_label_data.input_path
|
| self._pseudo_label_dataset_fn = pseudo_label_dataset_fn
|
| self._pseudo_label_data_ratio = params.pseudo_label_data.data_ratio
|
| self._pseudo_label_batch_size = params.pseudo_label_data.global_batch_size
|
| self._pseudo_label_matched_files = input_reader.match_files(
|
| self._pseudo_label_file_pattern)
|
| if not self._drop_remainder:
|
| raise ValueError(
|
| 'Must use drop_remainder=True with CombinationDatasetInputReader')
|
|
|
| def read(
|
| self,
|
| input_context: Optional[tf.distribute.InputContext] = None
|
| ) -> tf.data.Dataset:
|
| """Generates a tf.data.Dataset object."""
|
|
|
| labeled_batch_size, pl_batch_size = calculate_batch_sizes(
|
| self._global_batch_size, self._pseudo_label_data_ratio,
|
| self._pseudo_label_batch_size)
|
|
|
| if not labeled_batch_size and pl_batch_size:
|
| raise ValueError(
|
| 'Invalid batch_size: {} and pseudo_label_data_ratio: {}, '
|
| 'resulting in a 0 batch size for one of the datasets.'.format(
|
| self._global_batch_size, self._pseudo_label_data_ratio))
|
|
|
| def _read_decode_and_parse_dataset(matched_files, dataset_fn, batch_size,
|
| input_context):
|
| dataset = self._read_data_source(matched_files, dataset_fn, input_context)
|
| return self._decode_and_parse_dataset(dataset, batch_size, input_context)
|
|
|
| labeled_dataset = _read_decode_and_parse_dataset(
|
| matched_files=self._matched_files,
|
| dataset_fn=self._dataset_fn,
|
| batch_size=labeled_batch_size,
|
| input_context=input_context)
|
|
|
| pseudo_labeled_dataset = _read_decode_and_parse_dataset(
|
| matched_files=self._pseudo_label_matched_files,
|
| dataset_fn=self._pseudo_label_dataset_fn,
|
| batch_size=pl_batch_size,
|
| input_context=input_context)
|
|
|
| def concat_fn(d1, d2):
|
| return tf.nest.map_structure(
|
| lambda x1, x2: tf.concat([x1, x2], axis=0), d1, d2)
|
|
|
| dataset_concat = tf.data.Dataset.zip(
|
| (labeled_dataset, pseudo_labeled_dataset))
|
| dataset_concat = dataset_concat.map(
|
| concat_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
|
|
|
| def maybe_map_fn(dataset, fn):
|
| return dataset if fn is None else dataset.map(
|
| fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
|
|
|
| dataset_concat = maybe_map_fn(dataset_concat, self._postprocess_fn)
|
| dataset_concat = self._maybe_apply_data_service(dataset_concat,
|
| input_context)
|
|
|
| if self._deterministic is not None:
|
| options = tf.data.Options()
|
| options.experimental_deterministic = self._deterministic
|
| dataset_concat = dataset_concat.with_options(options)
|
|
|
| return dataset_concat.prefetch(tf.data.experimental.AUTOTUNE)
|
|
|
