official/vision/dataloaders/input_reader.py

# 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.

"""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:
    # Combine multiple datasets using weighted sampling.
    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)