Method load_dataset

tensorflow_datasets/core/load.py:302–383

Loads the named dataset from a dataset collection by calling `tfds.load`. Args: dataset: `str`, the dataset name to load. split: which split(s) of the dataset to load. If `None`, will return all splits available for the dataset. loader_kwargs: `dict` (optional), keywor

(
      self,
      dataset: str,
      split: Optional[Tree[splits_lib.SplitArg]] = None,
      loader_kwargs: dict[str, Any] | None = None,
  )

    self.loader_kwargs = loader_kwargs

  def load_dataset(
      self,
      dataset: str,
      split: Optional[Tree[splits_lib.SplitArg]] = None,
      loader_kwargs: dict[str, Any] | None = None,
  ) -> Mapping[str, tf.data.Dataset]:
    """Loads the named dataset from a dataset collection by calling `tfds.load`.

    Args:
      dataset: `str`, the dataset name to load.
      split: which split(s) of the dataset to load. If `None`, will return all
        splits available for the dataset.
      loader_kwargs: `dict` (optional), keyword arguments to be passed to the
        `tfds.load` function. Refer to `tfds.load` documentation for a
        comperehensive overview of the different loading options.

    Returns:
      A `dict` of {`str`: tf.data.Dataset} for the desided dataset.

    Raises:
      KeyError: if trying to load a dataset not included in the collection.
      RuntimeError: if `load` return type is not a `dict` or a `list`.
    """
    if not dataset:
      raise TypeError('You must specify a non-empty dataset to load.')

    loader_kwargs = loader_kwargs or self.loader_kwargs or {}

    # with_info must be False (or it will change the return type of `tfds.load`)
    if 'with_info' in loader_kwargs and loader_kwargs['with_info']:
      logging.warning('`with_info` cannot be True, setting it to False')
    loader_kwargs['with_info'] = False

    try:
      dataset_reference = self.datasets[dataset]
    except KeyError as e:
      raise KeyError(
          f'Dataset {dataset} is not included in this collection. '
          f'{self.collection.list_datasets(version=self.requested_version)}'
      ) from e

    # If `split` is defined both as argument and in `loader_kwargs`, always keep
    # the one defined as argument.
    if split:
      loader_kwargs['split'] = dataset_reference.get_split(split)
    # Make sure we always return a dict of dicts.
    if 'split' in loader_kwargs and isinstance(loader_kwargs['split'], str):
      loader_kwargs['split'] = [loader_kwargs['split']]

    # Add the data dir from the reference to loader_kwargs if it is defined and
    # not overridden in loader_kwargs.
    if (
        dataset_reference.data_dir is not None
        and 'data_dir' not in loader_kwargs
    ):
      loader_kwargs['data_dir'] = dataset_reference.data_dir

    load_output = load(dataset_reference.tfds_name(), **loader_kwargs)