### Add Dataset

# Writing custom datasets

Follow this guide to create a new dataset (either in TFDS or in your own
repository).

Check our [list of datasets](catalog/overview.md) to see if the dataset you want
is already present.

## TL;DR

The easiest way to write a new dataset is to use the
[TFDS CLI](https://www.tensorflow.org/datasets/cli):

```sh
cd path/to/my/project/datasets/
tfds new my_dataset  # Create `my_dataset/my_dataset.py` template files
# [...] Manually modify `my_dataset/my_dataset_dataset_builder.py` to implement your dataset.
cd my_dataset/
tfds build  # Download and prepare the dataset to `~/tensorflow_datasets/`
```

To use the new dataset with `tfds.load('my_dataset')`:

*   `tfds.load` will automatically detect and load the dataset generated in
    `~/tensorflow_datasets/my_dataset/` (e.g. by `tfds build`).
*   Alternatively, you can explicitly `import my.project.datasets.my_dataset` to
    register your dataset:

```python
import my.project.datasets.my_dataset  # Register `my_dataset`

ds = tfds.load('my_dataset')  # `my_dataset` registered
```

## Overview

Datasets are distributed in all kinds of formats and in all kinds of places, and
they're not always stored in a format that's ready to feed into a machine
learning pipeline. Enter TFDS.

TFDS process those datasets into a standard format (external data -> serialized
files), which can then be loaded as machine learning pipeline (serialized files
-> `tf.data.Dataset`). The serialization is done only once. Subsequent access
will read from those pre-processed files directly.

Most of the preprocessing is done automatically. Each dataset implements a
subclass of `tfds.core.DatasetBuilder`, which specifies:

*   Where the data is coming from (i.e. its URLs);
*   What the dataset looks like (i.e. its features);
*   How the data should be split (e.g. `TRAIN` and `TEST`);
*   and the individual examples in the dataset.

## Write your dataset

### Default template: `tfds new`

Use [TFDS CLI](https://www.tensorflow.org/datasets/cli) to generate the required
template python files.

```sh
cd path/to/project/datasets/  # Or use `--dir=path/to/project/datasets/` below
tfds new my_dataset
```

This command will generate a new `my_dataset/` folder with the following
structure:

```sh
my_dataset/
    __init__.py
    README.md # Markdown description of the dataset.
    CITATIONS.bib # Bibtex citation for the dataset.
    TAGS.txt # List of tags describing the dataset.
    my_dataset_dataset_builder.py # Dataset definition
    my_dataset_dataset_builder_test.py # Test
    dummy_data/ # (optional) Fake data (used for testing)
    checksum.tsv # (optional) URL checksums (see `checksums` section).
```

Search for `TODO(my_dataset)` here and modify accordingly.

### Dataset example

All datasets are implemented subclasses of `tfds.core.DatasetBuilder`, which
takes care of most boilerplate. It supports:

*   Small/medium datasets which can be generated on a single machine (this
    tutorial).
*   Very large datasets which require distributed generation (using
    [Apache Beam](https://beam.apache.org/), see our
    [huge dataset guide](https://www.tensorflow.org/datasets/beam_datasets#implementing_a_beam_dataset))

Here is a minimal example of a dataset builder that is based on
`tfds.core.GeneratorBasedBuilder`:

```python
class Builder(tfds.core.GeneratorBasedBuilder):
  """DatasetBuilder for my_dataset dataset."""

  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {
      '1.0.0': 'Initial release.',
  }

  def _info(self) -> tfds.core.DatasetInfo:
    """Dataset metadata (homepage, citation,...)."""
    return self.dataset_info_from_configs(
        features=tfds.features.FeaturesDict({
            'image': tfds.features.Image(shape=(256, 256, 3)),
            'label': tfds.features.ClassLabel(
                names=['no', 'yes'],
                doc='Whether this is a picture of a cat'),
        }),
    )

  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
    """Download the data and define splits."""
    extracted_path = dl_manager.download_and_extract('http://data.org/data.zip')
    # dl_manager returns pathlib-like objects with `path.read_text()`,
    # `path.iterdir()`,...
    return {
        'train': self._generate_examples(path=extracted_path / 'train_images'),
        'test': self._generate_examples(path=extracted_path / 'test_images'),
    }

  def _generate_examples(self, path) -> Iterator[Tuple[Key, Example]]:
    """Generator of examples for each split."""
    for img_path in path.glob('*.jpeg'):
      # Yields (key, example)
      yield img_path.name, {
          'image': img_path,
          'label': 'yes' if img_path.name.startswith('yes_') else 'no',
      }
```

Note that, for some specific data formats, we provide ready-to-use
[dataset builders](https://www.tensorflow.org/datasets/format_specific_dataset_builders)
to take care of most data processing.

Let's see in detail the 3 abstract methods to overwrite.

### `_info`: dataset metadata

`_info` returns the `tfds.core.DatasetInfo` containing the
[dataset metadata](https://www.tensorflow.org/datasets/overview#access_the_dataset_metadata).

```python
def _info(self):
  # The `dataset_info_from_configs` base method will construct the
  # `tfds.core.DatasetInfo` object using the passed-in parameters and
  # adding: builder (self), description/citations/tags from the config
  # files located in the same package.
  return self.dataset_info_from_configs(
      homepage='https://dataset-homepage.org',
      features=tfds.features.FeaturesDict({
          'image_description': tfds.features.Text(),
          'image': tfds.features.Image(),
          # Here, 'label' can be 0-4.
          'label': tfds.features.ClassLabel(num_classes=5),
      }),
      # If there's a common `(input, target)` tuple from the features,
      # specify them here. They'll be used if as_supervised=True in
      # builder.as_dataset.
      supervised_keys=('image', 'label'),
      # Specify whether to disable shuffling on the examples. Set to False by default.
      disable_shuffling=False,
  )
```

Most fields should be self-explanatory. Some precisions:

*   `features`: This specify the dataset structure, shape,... Support complex
    data types (audio, video, nested sequences,...). See the
    [available features](https://www.tensorflow.org/datasets/api_docs/python/tfds/features#classes)
    or the
    [feature connector guide](https://www.tensorflow.org/datasets/features) for
    more info.
*   `disable_shuffling`: See section
    [Maintain dataset order](#maintain-dataset-order).

Writing the `BibText` `CITATIONS.bib` file:

*   Search the dataset website for citation instruction (use that in BibTex
    format).
*   For [arXiv](https://arxiv.org/) papers: find the paper and click the
    `BibText` link on the right-hand side.
*   Find the paper on [Google Scholar](https://scholar.google.com) and click the
    double-quotation mark underneath the title and on the popup, click `BibTeX`.
*   If there is no associated paper (for example, there's just a website), you
    can use the [BibTeX Online Editor](https://truben.no/latex/bibtex/) to
    create a custom BibTeX entry (the drop-down menu has an `Online` entry
    type).

Updating the `TAGS.txt` file:

*   All allowed tags are pre-filled in the generated file.
*   Remove all tags which do not apply to the dataset.
*   Valid tags are listed in
    [tensorflow_datasets/core/valid_tags.txt](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/core/valid_tags.txt).
*   To add a tag to that list, please send a PR.

#### Maintain dataset order

By default, the records of the datasets are shuffled when stored in order to
make the distribution of classes more uniform across the dataset, since often
records belonging to the same class are contiguous. In order to specify that the
dataset should be sorted by the key generated provided by `_generate_examples`
the field `disable_shuffling` should be set to `True`. By default it is set to
`False`.

```python
def _info(self):
  return self.dataset_info_from_configs(
    # [...]
    disable_shuffling=True,
    # [...]
  )
```

Keep in mind that disabling shuffling has a performance impact as shards cannot
be read in parallel anymore.

### `_split_generators`: downloads and splits data

#### Downloading and extracting source data

Most datasets need to download data from the web. This is done using the
`tfds.download.DownloadManager` input argument of `_split_generators`.
`dl_manager` has the following methods:

*   `download`: supports `http(s)://`, `ftp(s)://`
*   `extract`: currently supports `.zip`, `.gz`, and `.tar` files.
*   `download_and_extract`: Same as
    `dl_manager.extract(dl_manager.download(urls))`

All those methods returns `tfds.core.Path` (aliases for
[`epath.Path`](https://github.com/google/etils)), which are
[pathlib.Path-like](https://docs.python.org/3/library/pathlib.html) objects.

Those methods supports arbitrary nested structure (`list`, `dict`), like:

```python
extracted_paths = dl_manager.download_and_extract({
    'foo': 'https://example.com/foo.zip',
    'bar': 'https://example.com/bar.zip',
})
# This returns:
assert extracted_paths == {
    'foo': Path('/path/to/extracted_foo/'),
    'bar': Path('/path/extracted_bar/'),
}
```

#### Manual download and extraction

Some data cannot be automatically downloaded (e.g. require a login), in this
case, user will manually download the source data and place it in `manual_dir/`
(defaults to `~/tensorflow_datasets/downloads/manual/`).

Files can then be accessed through `dl_manager.manual_dir`:

```python
class MyDataset(tfds.core.GeneratorBasedBuilder):

  MANUAL_DOWNLOAD_INSTRUCTIONS = """
  Register into https://example.org/login to get the data. Place the `data.zip`
  file in the `manual_dir/`.
  """

  def _split_generators(self, dl_manager):
    # data_path is a pathlib-like `Path('<manual_dir>/data.zip')`
    archive_path = dl_manager.manual_dir / 'data.zip'
    # Extract the manually downloaded `data.zip`
    extracted_path = dl_manager.extract(archive_path)
    ...
```

The `manual_dir` location can be customized with `tfds build --manual_dir=` or
using `tfds.download.DownloadConfig`.

#### Read archive directly

`dl_manager.iter_archive` reads an archives sequentially without extracting
them. This can save storage space and improve performances on some file systems.

```python
for filename, fobj in dl_manager.iter_archive('path/to/archive.zip'):
  ...
```

`fobj` has the same methods as `with open('rb') as fobj:` (e.g. `fobj.read()`)

#### Specifying dataset splits

If the dataset comes with pre-defined splits (e.g. `MNIST` has `train` and
`test` splits), keep those. Otherwise, only specify a single `all` split. Users
can dynamically create their own subsplits with the
[subsplit API](https://www.tensorflow.org/datasets/splits) (e.g.
`split='train[80%:]'`). Note that any alphabetical string can be used as split
name, apart from the aforementioned `all`.

```python
def _split_generators(self, dl_manager):
  # Download source data
  extracted_path = dl_manager.download_and_extract(...)

  # Specify the splits
  return {
      'train': self._generate_examples(
          images_path=extracted_path / 'train_imgs',
          label_path=extracted_path / 'train_labels.csv',
      ),
      'test': self._generate_examples(
          images_path=extracted_path / 'test_imgs',
          label_path=extracted_path / 'test_labels.csv',
      ),
  }
```

### `_generate_examples`: Example generator

`_generate_examples` generates the examples for each split from the source data.

This method will typically read source dataset artifacts (e.g. a CSV file) and
yield `(key, feature_dict)` tuples:

*   `key`: Example identifier. Used to deterministically shuffle the examples
    using `hash(key)` or to sort by key when shuffling is disabled (see section
    [Maintain dataset order](#maintain-dataset-order)). Should be:
    *   **unique**: If two examples use the same key, an exception will be
        raised.
    *   **deterministic**: Should not depend on `download_dir`,
        `os.path.listdir` order,... Generating the data twice should yield the
        same key.
    *   **comparable**: If shuffling is disabled the key will be used to sort
        the dataset.
*   `feature_dict`: A `dict` containing the example values.
    *   The structure should match the `features=` structure defined in
        `tfds.core.DatasetInfo`.
    *   Complex data types (image, video, audio,...) will be automatically
        encoded.
    *   Each feature often accept multiple input types (e.g. video accept
        `/path/to/vid.mp4`, `np.array(shape=(l, h, w, c))`, `List[paths]`,
        `List[np.array(shape=(h, w, c)]`, `List[img_bytes]`,...)
    *   See the
        [feature connector guide](https://www.tensorflow.org/datasets/features)
        for more info.

```python
def _generate_examples(self, images_path, label_path):
  # Read the input data out of the source files
  with label_path.open() as f:
    for row in csv.DictReader(f):
      image_id = row['image_id']
      # And yield (key, feature_dict)
      yield image_id, {
          'image_description': row['description'],
          'image': images_path / f'{image_id}.jpeg',
          'label': row['label'],
      }
```

Warning: When parsing boolean values from strings or integers, use the util
function `tfds.core.utils.bool_utils.parse_bool` to avoid parsing errors (e.g.,
`bool("False") == True`).

#### File access and `tf.io.gfile`

In order to support Cloud storage systems, avoid the use of the Python built-in
I/O ops.

Instead, the `dl_manager` returns
[pathlib-like](https://docs.python.org/3/library/pathlib.html) objects directly
compatible with Google Cloud storage:

```python
path = dl_manager.download_and_extract('http://some-website/my_data.zip')

json_path = path / 'data/file.json'

json.loads(json_path.read_text())
```

Alternatively, use `tf.io.gfile` API instead of built-in for file operations:

*   `open` -> `tf.io.gfile.GFile`
*   `os.rename` -> `tf.io.gfile.rename`
*   ...

Pathlib should be prefered to `tf.io.gfile` (see
[rational](https://www.tensorflow.org/datasets/common_gotchas#prefer_to_use_pathlib_api).

#### Extra dependencies

Some datasets require additional Python dependencies only during generation. For
example, the SVHN dataset uses `scipy` to load some data.

If you're adding dataset into the TFDS repository, please use
`tfds.core.lazy_imports` to keep the `tensorflow-datasets` package small. Users
will install additional dependencies only as needed.

To use `lazy_imports`:

*   Add an entry for your dataset into `DATASET_EXTRAS` in
    [`setup.py`](https://github.com/tensorflow/datasets/tree/master/setup.py).
    This makes it so that users can do, for example, `pip install
    'tensorflow-datasets[svhn]'` to install the extra dependencies.
*   Add an entry for your import to
    [`LazyImporter`](https://github.com/tensorflow/datasets/tree/master/tensorflow_datasets/core/lazy_imports_lib.py)
    and to the
    [`LazyImportsTest`](https://github.com/tensorflow/datasets/tree/master/tensorflow_datasets/core/lazy_imports_lib_test.py).
*   Use `tfds.core.lazy_imports` to access the dependency (for example,
    `tfds.core.lazy_imports.scipy`) in your `DatasetBuilder`.

#### Corrupted data

Some datasets are not perfectly clean and contain some corrupt data (for
example, the images are in JPEG files but some are invalid JPEG). These examples
should be skipped, but leave a note in the dataset description how many examples
were dropped and why.

### Dataset configuration/variants (tfds.core.BuilderConfig)

Some datasets may have multiple variants, or options for how the data is
preprocessed and written to disk. For example,
[cycle_gan](https://www.tensorflow.org/datasets/catalog/cycle_gan) has one
config per object pairs (`cycle_gan/horse2zebra`, `cycle_gan/monet2photo`,...).

This is done through `tfds.core.BuilderConfig`s:

1.  Define your configuration object as a subclass of `tfds.core.BuilderConfig`.
    For example, `MyDatasetConfig`.

    ```python
    @dataclasses.dataclass
    class MyDatasetConfig(tfds.core.BuilderConfig):
      img_size: Tuple[int, int] = (0, 0)
    ```

    Note: Default values are required because of
    https://bugs.python.org/issue33129.

1.  Define the `BUILDER_CONFIGS = []` class member in `MyDataset` that lists
    `MyDatasetConfig`s that the dataset exposes.

    ```python
    class MyDataset(tfds.core.GeneratorBasedBuilder):
      VERSION = tfds.core.Version('1.0.0')
      # pytype: disable=wrong-keyword-args
      BUILDER_CONFIGS = [
          # `name` (and optionally `description`) are required for each config
          MyDatasetConfig(name='small', description='Small ...', img_size=(8, 8)),
          MyDatasetConfig(name='big', description='Big ...', img_size=(32, 32)),
      ]
      # pytype: enable=wrong-keyword-args
    ```

    Note: `# pytype: disable=wrong-keyword-args` is required because of
    [Pytype bug](https://github.com/google/pytype/issues/628) with dataclasses
    inheritance.

1.  Use `self.builder_config` in `MyDataset` to configure data generation (e.g.
    `shape=self.builder_config.img_size`). This may include setting different
    values in `_info()` or changing download data access.

Notes:

*   Each config has a unique name. The fully qualified name of a config is
    `dataset_name/config_name` (e.g. `coco/2017`).
*   If not specified, the first config in `BUILDER_CONFIGS` will be used (e.g.
    `tfds.load('c4')` default to `c4/en`)

See
[`anli`](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/text/anli.py#L69)
for an example of a dataset that uses `BuilderConfig`s.

### Version

Version can refer to two different meaning:

*   The "external" original data version: e.g. COCO v2019, v2017,...
*   The "internal" TFDS code version: e.g. rename a feature in
    `tfds.features.FeaturesDict`, fix a bug in `_generate_examples`

To update a dataset:

*   For "external" data update: Multiple users may want to access a specific
    year/version simultaneously. This is done by using one
    `tfds.core.BuilderConfig` per version (e.g. `coco/2017`, `coco/2019`) or one
    class per version (e.g. `Voc2007`, `Voc2012`).
*   For "internal" code update: Users only download the most recent version. Any
    code update should increase the `VERSION` class attribute (e.g. from `1.0.0`
    to `VERSION = tfds.core.Version('2.0.0')`) following
    [semantic versioning](https://www.tensorflow.org/datasets/datasets_versioning#semantic).

### Add an import for registration

Don't forget to import the dataset module to your project `__init__` to be
automatically registered in `tfds.load`, `tfds.builder`.

```python
import my_project.datasets.my_dataset  # Register MyDataset

ds = tfds.load('my_dataset')  # MyDataset available
```

For example, if you're contributing to `tensorflow/datasets`, add the module
import to its subdirectory's `__init__.py` (e.g.
[`image/__init__.py`](https://github.com/tensorflow/datasets/tree/master/tensorflow_datasets/image/__init__.py).

### Check for common implementation gotchas

Please check for the
[common implementation gotchas](https://www.tensorflow.org/datasets/common_gotchas).

## Test your dataset

### Download and prepare: `tfds build`

To generate the dataset, run `tfds build` from the `my_dataset/` directory:

```sh
cd path/to/datasets/my_dataset/
tfds build --register_checksums
```

Some useful flags for development:

*   `--pdb`: Enter debugging mode if an exception is raised.
*   `--overwrite`: Delete existing files if the dataset was already generated.
*   `--max_examples_per_split`: Only generate the first X examples (default to
    1), rather than the full dataset.
*   `--register_checksums`: Record the checksums of downloaded urls. Should only
    be used while in development.

See the
[CLI documentation](https://www.tensorflow.org/datasets/cli#tfds_build_download_and_prepare_a_dataset)
for full list of flags.

### Checksums

It is recommended to record the checksums of your datasets to guarantee
determinism, help with documentation,... This is done by generating the dataset
with the `--register_checksums` (see previous section).

If you are releasing your datasets through PyPI, don't forget to export the
`checksums.tsv` files (e.g. in the `package_data` of your `setup.py`).

### Unit-test your dataset

`tfds.testing.DatasetBuilderTestCase` is a base `TestCase` to fully exercise a
dataset. It uses "dummy data" as test data that mimic the structure of the
source dataset.

*   The test data should be put in `my_dataset/dummy_data/` directory and should
    mimic the source dataset artifacts as downloaded and extracted. It can be
    created manually or automatically with a script
    ([example script](https://github.com/tensorflow/datasets/tree/master/tensorflow_datasets/datasets/bccd/dummy_data_generation.py)).
*   Make sure to use different data in your test data splits, as the test will
    fail if your dataset splits overlap.
*   **The test data should not contain any copyrighted material**. If in doubt,
    do not create the data using material from the original dataset.

```python
import tensorflow_datasets as tfds
from . import my_dataset_dataset_builder


class MyDatasetTest(tfds.testing.DatasetBuilderTestCase):
  """Tests for my_dataset dataset."""
  DATASET_CLASS = my_dataset_dataset_builder.Builder
  SPLITS = {
      'train': 3,  # Number of fake train example
      'test': 1,  # Number of fake test example
  }

  # If you are calling `download/download_and_extract` with a dict, like:
  #   dl_manager.download({'some_key': 'http://a.org/out.txt', ...})
  # then the tests needs to provide the fake output paths relative to the
  # fake data directory
  DL_EXTRACT_RESULT = {
      'name1': 'path/to/file1',  # Relative to my_dataset/dummy_data dir.
      'name2': 'file2',
  }


if __name__ == '__main__':
  tfds.testing.test_main()
```

Run the following command to test the dataset.

```sh
python my_dataset_test.py
```

## Send us feedback

We are continuously trying to improve the dataset creation workflow, but can
only do so if we are aware of the issues. Which issues or errors did you
encounter while creating the dataset? Was there a part which was confusing, or
wasn't working the first time?

Please share your feedback on
[GitHub](https://github.com/tensorflow/datasets/issues).

---

### Add Dataset Collection

# Add a new dataset collection

Follow this guide to create a new dataset collection (either in TFDS or in your
own repository).

## Overview

To add a new dataset collection `my_collection` to TFDS, users need to generate
a `my_collection` folder containing the following files:

```sh
my_collection/
  __init__.py
  my_collection.py # Dataset collection definition
  my_collection_test.py # (Optional) test
  description.md # (Optional) collection description (if not included in my_collection.py)
  citations.md # (Optional) collection citations (if not included in my_collection.py)
```

As a convention, new dataset collections should be added to the
`tensorflow_datasets/dataset_collections/` folder in the TFDS repository.

## Write your dataset collection

All dataset collections are implemented subclasses of
[`tfds.core.dataset_collection_builder.DatasetCollection`](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/core/dataset_collection_builder.py).

Here is a minimal example of a dataset collection builder, defined in the file
`my_collection.py`:

```python
import collections
from typing import Mapping
from tensorflow_datasets.core import dataset_collection_builder
from tensorflow_datasets.core import naming

class MyCollection(dataset_collection_builder.DatasetCollection):
  """Dataset collection builder my_dataset_collection."""

  @property
  def info(self) -> dataset_collection_builder.DatasetCollectionInfo:
    return dataset_collection_builder.DatasetCollectionInfo.from_cls(
        dataset_collection_class=self.__class__,
        description="my_dataset_collection description.",
        release_notes={
            "1.0.0": "Initial release",
        },
    )

  @property
  def datasets(
      self,
  ) -> Mapping[str, Mapping[str, naming.DatasetReference]]:
    return collections.OrderedDict({
        "1.0.0":
            naming.references_for({
                "dataset_1": "natural_questions/default:0.0.2",
                "dataset_2": "media_sum:1.0.0",
            }),
        "1.1.0":
            naming.references_for({
                "dataset_1": "natural_questions/longt5:0.1.0",
                "dataset_2": "media_sum:1.0.0",
                "dataset_3": "squad:3.0.0"
            })
    })
```

The next sections describe the 2 abstract methods to overwrite.

### `info`: dataset collection metadata

The `info` method returns the
[`dataset_collection_builder.DatasetCollectionInfo`](https://github.com/tensorflow/datasets/blob/4854e55ddf5fb68c63ddbd502ad0ef4ec6e08b40/tensorflow_datasets/core/dataset_collection_builder.py#L66)
containing the collection's metadata.

The dataset collection info contains four fields:

*   name: the name of the dataset collection.
*   description: a markdown-formatted description of the dataset collection.
    There are two ways to define a dataset collection's description: (1) As a
    (multi-line) string directly in the collection's `my_collection.py` file -
    similarly as it is already done for TFDS datasets; (2) In a `description.md`
    file, which must be placed in the dataset collection folder.
*   release_notes: a mapping from the dataset collection's version to the
    corresponding release notes.
*   citation: An optional (list of) BibTeX citation(s) for the dataset
    collection. There are two ways to define a dataset collection's citation:
    (1) As a (multi-line) string directly in the collection's `my_collection.py`
    file - similarly as it is already done for TFDS datasets; (2) In a
    `citations.bib` file, which must be placed in the dataset collection folder.

### `datasets`: define the datasets in the collection

The `datasets` method returns the TFDS datasets in the collection.

It is defined as a dictionary of versions, which describe the evolution of the
dataset collection.

For each version, the included TFDS datasets are stored as a dictionary from
dataset names to
[`naming.DatasetReference`](https://github.com/tensorflow/datasets/blob/4854e55ddf5fb68c63ddbd502ad0ef4ec6e08b40/tensorflow_datasets/core/naming.py#L187).
For example:

```python
class MyCollection(dataset_collection_builder.DatasetCollection):
  ...
  @property
  def datasets(self):
    return {
        "1.0.0": {
            "yes_no":
                naming.DatasetReference(
                    dataset_name="yes_no", version="1.0.0"),
            "sst2":
                naming.DatasetReference(
                    dataset_name="glue", config="sst2", version="2.0.0"),
            "assin2":
                naming.DatasetReference(
                    dataset_name="assin2", version="1.0.0"),
        },
        ...
    }
```

The
[`naming.references_for`](https://github.com/tensorflow/datasets/blob/4854e55ddf5fb68c63ddbd502ad0ef4ec6e08b40/tensorflow_datasets/core/naming.py#L257)
method provides a more compact way to express the same as above:

```python
class MyCollection(dataset_collection_builder.DatasetCollection):
  ...
  @property
  def datasets(self):
    return {
        "1.0.0":
            naming.references_for({
                "yes_no": "yes_no:1.0.0",
                "sst2": "glue/sst:2.0.0",
                "assin2": "assin2:1.0.0",
            }),
        ...
    }
```

## Unit-test your dataset collection

[DatasetCollectionTestBase](https://github.com/tensorflow/datasets/blob/4854e55ddf5fb68c63ddbd502ad0ef4ec6e08b40/tensorflow_datasets/testing/dataset_collection_builder_testing.py#L28)
is a base test class for dataset collections. It provides a number of simple
checks to guarantee that the dataset collection is correctly registered, and its
datasets exist in TFDS.

The only class attribute to set is `DATASET_COLLECTION_CLASS`, which specifies
the class object of dataset collection to test.

Additionally, users can set the following class attributes:

*   `VERSION`: The version of the dataset collection used to run the test
    (defaults to the latest version).
*   `DATASETS_TO_TEST`: List containing the datasets to test existence for in
    TFDS (defaults to all datasets in the collection).
*   `CHECK_DATASETS_VERSION`: Whether to check for the existence of the
    versioned datasets in the dataset collection, or for their default versions
    (defaults to true).

The simplest valid test for a dataset collection would be:

```python
from tensorflow_datasets.testing.dataset_collection_builder_testing import DatasetCollectionTestBase
from . import my_collection

class TestMyCollection(DatasetCollectionTestBase):
  DATASET_COLLECTION_CLASS = my_collection.MyCollection
```

Run the following command to test the dataset collection.

```sh
python my_dataset_test.py
```

## Feedback

We are continuously trying to improve the dataset creation workflow, but can
only do so if we are aware of the issues. Which issues or errors did you
encounter while creating the dataset collection? Was there a part which was
confusing, or wasn't working the first time?

Please share your feedback on
[GitHub](https://github.com/tensorflow/datasets/issues).

---

### Announce Proxy

# Introducing TensorFlow Datasets

TensorFlow Datasets is now released on PyPI:

`pip install tensorflow-datasets`

Read the [blog post](https://medium.com/tensorflow/introducing-tensorflow-datasets-c7f01f7e19f3)
to learn more.

---

### Beam Datasets

# Generating big datasets with Apache Beam

Some datasets are too big to be processed on a single machine. `tfds` supports
generating data across many machines by using
[Apache Beam](https://beam.apache.org/).

This doc has two sections:

*   For user who want to generate an existing Beam dataset
*   For developers who want to create a new Beam dataset

## Generating a Beam dataset

Below are different examples of generating a Beam dataset, both on the cloud or
locally.

**Warning**: When generating the dataset with the
[`tfds build` CLI](https://www.tensorflow.org/datasets/cli#tfds_build_download_and_prepare_a_dataset),
make sure to specify the dataset config you want to generate or it will default
to generate all existing configs. For example, for
[wikipedia](https://www.tensorflow.org/datasets/catalog/wikipedia), use `tfds
build wikipedia/20200301.en` instead of `tfds build wikipedia`.

### On Google Cloud Dataflow

To run the pipeline using
[Google Cloud Dataflow](https://cloud.google.com/dataflow/) and take advantage
of distributed computation, first follow the
[Quickstart instructions](https://cloud.google.com/dataflow/docs/quickstarts/quickstart-python).

Once your environment is set up, you can run the
[`tfds build` CLI](https://www.tensorflow.org/datasets/cli#tfds_build_download_and_prepare_a_dataset)
using a data directory on [GCS](https://cloud.google.com/storage/) and
specifying the
[required options](https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#configuring-pipelineoptions-for-execution-on-the-cloud-dataflow-service)
for the `--beam_pipeline_options` flag.

To make it easier to launch the script, it's helpful to define the following
variables using the actual values for your GCP/GCS setup and the dataset you
want to generate:

```sh
DATASET_NAME=<dataset-name>
DATASET_CONFIG=<dataset-config>
GCP_PROJECT=my-project-id
GCS_BUCKET=gs://my-gcs-bucket
```

You will then need to create a file to tell Dataflow to install `tfds` on the
workers:

```sh
echo "tensorflow_datasets[$DATASET_NAME]" > /tmp/beam_requirements.txt
```

If you're using `tfds-nightly`, make sure to echo from `tfds-nightly` in case
the dataset has been updated since the last release.

```sh
echo "tfds-nightly[$DATASET_NAME]" > /tmp/beam_requirements.txt
```

If you're using additional dependencies not included in TFDS library follow
[the instructions for managing Python pipeline dependencies](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/#local-or-nonpypi).

Finally, you can launch the job using the command below:

```sh
tfds build $DATASET_NAME/$DATASET_CONFIG \
  --data_dir=$GCS_BUCKET/tensorflow_datasets \
  --beam_pipeline_options=\
"runner=DataflowRunner,project=$GCP_PROJECT,job_name=$DATASET_NAME-gen,"\
"staging_location=$GCS_BUCKET/binaries,temp_location=$GCS_BUCKET/temp,"\
"requirements_file=/tmp/beam_requirements.txt"
```

### Locally

To run your script locally using the
[default Apache Beam runner](https://beam.apache.org/documentation/runners/direct/)
(it must fit all data in memory), the command is the same as for other datasets:

```sh
tfds build my_dataset
```

**Warning**: Beam datasets can be **huge** (terabytes or larger) and take a
significant amount of resources to be generated (can take weeks on a local
computer). It is recommended to generate the datasets using a distributed
environment. Have a look at the
[Apache Beam Documentation](https://beam.apache.org/) for a list of supported
runtimes.

### With Apache Flink

To run the pipeline using [Apache Flink](https://flink.apache.org/) you can read
the
[official documentation](https://beam.apache.org/documentation/runners/flink).
Make sure your Beam is compliant with
[Flink Version Compatibility](https://beam.apache.org/documentation/runners/flink/#flink-version-compatibility)

To make it easier to launch the script, it's helpful to define the following
variables using the actual values for your Flink setup and the dataset you want
to generate:

```sh
DATASET_NAME=<dataset-name>
DATASET_CONFIG=<dataset-config>
FLINK_CONFIG_DIR=<flink-config-directory>
FLINK_VERSION=<flink-version>
```

To run on an embedded Flink cluster, you can launch the job using the command
below:

```sh
tfds build $DATASET_NAME/$DATASET_CONFIG \
  --beam_pipeline_options=\
"runner=FlinkRunner,flink_version=$FLINK_VERSION,flink_conf_dir=$FLINK_CONFIG_DIR"
```

### With a custom script

To generate the dataset on Beam, the API is the same as for other datasets. You
can customize the
[`beam.Pipeline`](https://beam.apache.org/documentation/programming-guide/#creating-a-pipeline)
using the `beam_options` (and `beam_runner`) arguments of `DownloadConfig`.

```python
# If you are running on Dataflow, Spark,..., you may have to set-up runtime
# flags. Otherwise, you can leave flags empty [].
flags = ['--runner=DataflowRunner', '--project=<project-name>', ...]

# `beam_options` (and `beam_runner`) will be forwarded to `beam.Pipeline`
dl_config = tfds.download.DownloadConfig(
    beam_options=beam.options.pipeline_options.PipelineOptions(flags=flags)
)
data_dir = 'gs://my-gcs-bucket/tensorflow_datasets'
builder = tfds.builder('wikipedia/20190301.en', data_dir=data_dir)
builder.download_and_prepare(download_config=dl_config)
```

## Implementing a Beam dataset

### Prerequisites

In order to write Apache Beam datasets, you should be familiar with the
following concepts:

*   Be familiar with the
    [`tfds` dataset creation guide](https://github.com/tensorflow/datasets/blob/master/docs/add_dataset.md)
    as most of the content still applies for Beam datasets.
*   Get an introduction to Apache Beam with the
    [Beam programming guide](https://beam.apache.org/documentation/programming-guide/).
*   If you want to generate your dataset using Cloud Dataflow, read the
    [Google Cloud Documentation](https://cloud.google.com/dataflow/docs/quickstarts/quickstart-python)
    and the
    [Apache Beam dependency guide](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/).

### Instructions

If you are familiar with the
[dataset creation guide](https://github.com/tensorflow/datasets/blob/master/docs/add_dataset.md),
adding a Beam dataset only requires to modify the `_generate_examples` function.
The function should returns a beam object, rather than a generator:

Non-beam dataset:

```python
def _generate_examples(self, path):
  for f in path.iterdir():
    yield _process_example(f)
```

Beam dataset:

```python
def _generate_examples(self, path):
  return (
      beam.Create(path.iterdir())
      | beam.Map(_process_example)
  )
```

All the rest can be 100% identical, including tests.

Some additional considerations:

*   Use `tfds.core.lazy_imports` to import Apache Beam. By using a lazy
    dependency, users can still read the dataset after it has been generated
    without having to install Beam.
*   Be careful with Python closures. When running the pipeline, the `beam.Map`
    and `beam.DoFn` functions are serialized using `pickle` and sent to all
    workers. Do not use mutable objects inside a `beam.PTransform` if the state
    has to be shared across workers.
*   Due to the way `tfds.core.DatasetBuilder` is serialized with pickle,
    mutating `tfds.core.DatasetBuilder` during data creation will be ignored on
    the workers (e.g. it's not possible to set `self.info.metadata['offset'] =
    123` in `_split_generators` and access it from the workers like
    `beam.Map(lambda x: x + self.info.metadata['offset'])`)
*   If you need to share some pipeline steps between the splits, you can add add
    an extra `pipeline: beam.Pipeline` kwarg to `_split_generator` and control
    the full generation pipeline. See `_generate_examples` documentation of
    `tfds.core.GeneratorBasedBuilder`.

### Example

Here is an example of a Beam dataset.

```python
class DummyBeamDataset(tfds.core.GeneratorBasedBuilder):

  VERSION = tfds.core.Version('1.0.0')

  def _info(self):
    return self.dataset_info_from_configs(
        features=tfds.features.FeaturesDict({
            'image': tfds.features.Image(shape=(16, 16, 1)),
            'label': tfds.features.ClassLabel(names=['dog', 'cat']),
        }),
    )

  def _split_generators(self, dl_manager):
    ...
    return {
        'train': self._generate_examples(file_dir='path/to/train_data/'),
        'test': self._generate_examples(file_dir='path/to/test_data/'),
    }

  def _generate_examples(self, file_dir: str):
    """Generate examples as dicts."""
    beam = tfds.core.lazy_imports.apache_beam

    def _process_example(filename):
      # Use filename as key
      return filename, {
          'image': os.path.join(file_dir, filename),
          'label': filename.split('.')[1],  # Extract label: "0010102.dog.jpeg"
      }

    return (
        beam.Create(tf.io.gfile.listdir(file_dir))
        | beam.Map(_process_example)
    )

```

### Running your pipeline

To run the pipeline, have a look at the above section.

**Note**: Like for non-beam datasets, do not forget to register download
checksums with `--register_checksums` (only the first time to register the
downloads).

```sh
tfds build my_dataset --register_checksums
```

## Pipeline using TFDS as input

If you want to create a beam pipeline which takes a TFDS dataset as source, you
can use the `tfds.beam.ReadFromTFDS`:

```python
builder = tfds.builder('my_dataset')

_ = (
    pipeline
    | tfds.beam.ReadFromTFDS(builder, split='train')
    | beam.Map(tfds.as_numpy)
    | ...
)
```

It will process each shard of the dataset in parallel.

Note: This require the dataset to be already generated. To generate datasets
using beam, see the other sections.

---

### Common Gotchas

# Common implementation gotchas

This page describe the common implementation gotcha when implementing a new
dataset.

## Legacy `SplitGenerator` should be avoided

The old `tfds.core.SplitGenerator` API is deprecated.

```python
def _split_generator(...):
  return [
      tfds.core.SplitGenerator(name='train', gen_kwargs={'path': train_path}),
      tfds.core.SplitGenerator(name='test', gen_kwargs={'path': test_path}),
  ]
```

Should be replaced by:

```python
def _split_generator(...):
  return {
      'train': self._generate_examples(path=train_path),
      'test': self._generate_examples(path=test_path),
  }
```

**Rationale**: The new API is less verbose and more explicit. The old API will
be removed in future version.

## New datasets should be self-contained in a folder

When adding a dataset inside the `tensorflow_datasets/` repository, please make
sure to follow the dataset-as-folder structure (all checksums, dummy data,
implementation code self-contained in a folder).

*   Old datasets (bad): `<category>/<ds_name>.py`
*   New datasets (good): `<category>/<ds_name>/<ds_name>.py`

Use the
[TFDS CLI](https://www.tensorflow.org/datasets/cli#tfds_new_implementing_a_new_dataset)
(`tfds new`, or `gtfds new` for googlers) to generate the template.

**Rationale**: Old structure required absolute paths for checksums, fake data
and was distributing the dataset files in many places. It was making it harder
to implement datasets outside the TFDS repository. For consistency, the new
structure should be used everywhere now.

## Description lists should be formatted as markdown

The `DatasetInfo.description` `str` is formatted as markdown. Markdown lists
require an empty line before the first item:

```python
_DESCRIPTION = """
Some text.
                      # << Empty line here !!!
1. Item 1
2. Item 1
3. Item 1
                      # << Empty line here !!!
Some other text.
"""
```

**Rationale**: Badly formatted description create visual artifacts in our
catalog documentation. Without the empty lines, the above text would be rendered
as:

Some text. 1. Item 1 2. Item 1 3. Item 1 Some other text

## Forgot ClassLabel names

When using `tfds.features.ClassLabel`, try to provide the human-readable labels
`str` with `names=` or `names_file=` (instead of `num_classes=10`).

```python
features = {
    'label': tfds.features.ClassLabel(names=['dog', 'cat', ...]),
}
```

**Rationale**: Human readable labels are used in many places:

*   Allow to yield `str` directly in `_generate_examples`: `yield {'label':
    'dog'}`
*   Exposed in the users like `info.features['label'].names` (conversion method
    `.str2int('dog')`,... also available)
*   Used in the
    [visualization utils](https://www.tensorflow.org/datasets/overview#tfdsas_dataframe)
    `tfds.show_examples`, `tfds.as_dataframe`

## Forgot image shape

When using `tfds.features.Image`, `tfds.features.Video`, if the images have
static shape, they should be explicitly specified:

```python
features = {
    'image': tfds.features.Image(shape=(256, 256, 3)),
}
```

**Rationale**: It allow static shape inference (e.g.
`ds.element_spec['image'].shape`), which is required for batching (batching
images of unknown shape would require resizing them first).

## Prefer more specific type instead of `tfds.features.Tensor`

When possible, prefer the more specific types `tfds.features.ClassLabel`,
`tfds.features.BBoxFeatures`,... instead of the generic `tfds.features.Tensor`.

**Rationale**: In addition of being more semantically correct, specific features
provides additional metadata to users and are detected by tools.

## Lazy imports in global space

Lazy imports should not be called from the global space. For example the
following is wrong:

```python
tfds.lazy_imports.apache_beam # << Error: Import beam in the global scope

def f() -> beam.Map:
  ...
```

**Rationale**: Using lazy imports in the global scope would import the module
for all tfds users, defeating the purpose of lazy imports.

## Dynamically computing train/test splits

If the dataset does not provide official splits, neither should TFDS. The
following should be avoided:

```python
_TRAIN_TEST_RATIO = 0.7

def _split_generator():
  ids = list(range(num_examples))
  np.random.RandomState(seed).shuffle(ids)

  # Split train/test
  train_ids = ids[_TRAIN_TEST_RATIO * num_examples:]
  test_ids = ids[:_TRAIN_TEST_RATIO * num_examples]
  return {
      'train': self._generate_examples(train_ids),
      'test': self._generate_examples(test_ids),
  }
```

**Rationale**: TFDS try to provide datasets as close as the original data. The
[sub-split API](https://www.tensorflow.org/datasets/splits) should be used
instead to let users dynamically create the subsplits they want:

```python
ds_train, ds_test = tfds.load(..., split=['train[:80%]', 'train[80%:]'])
```

## Python style guide

### Prefer to use pathlib API

Instead of the `tf.io.gfile` API, it is preferable to use the
[pathlib API](https://docs.python.org/3/library/pathlib.html). All `dl_manager`
methods returns pathlib-like objects compatible with GCS, S3,...

```python
path = dl_manager.download_and_extract('http://some-website/my_data.zip')

json_path = path / 'data/file.json'

json.loads(json_path.read_text())
```

**Rationale**: pathlib API is a modern object oriented file API which remove
boilerplate. Using `.read_text()` / `.read_bytes()` also guarantee the files are
correctly closed.

### If the method is not using `self`, it should be a function

If a class method is not using `self`, it should be a simple function (defined
outside the class).

**Rationale**: It makes it explicit to the reader that the function do not have
side effects, nor hidden input/output:

```python
x = f(y)  # Clear inputs/outputs

x = self.f(y)  # Does f depend on additional hidden variables ? Is it stateful ?
```

## Lazy imports in Python

We lazily import big modules like TensorFlow. Lazy imports defer the actual
import of the module to the first usage of the module. So users who don't need
this big module will never import it. We use `etils.epy.lazy_imports`.

```python
from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
# After this statement, TensorFlow is not imported yet

...

features = tfds.features.Image(dtype=tf.uint8)
# After using it (`tf.uint8`), TensorFlow is now imported
```

Under the hood, the
[`LazyModule` class](https://github.com/google/etils/blob/main/etils/epy/lazy_imports_utils.py)
acts as a factory, that will only actually import the module when an attribute
is accessed (`__getattr__`).

You can also use it conveniently with a context manager:

```python
from etils import epy

with epy.lazy_imports(error_callback=..., success_callback=...):
  import some_big_module
```

---

### Contribute

# Contribute to the TFDS repository

Thank you for your interest in our library ! We are thrilled to have such a
motivated community.

## Get started

*   If you're new with TFDS, the easiest way to get started is to implement one
    of our
    [requested datasets](https://github.com/tensorflow/datasets/issues?q=is%3Aissue+is%3Aopen+label%3A%22dataset+request%22+sort%3Areactions-%2B1-desc),
    focusing on the most requested ones.
    [Follow our guide](https://www.tensorflow.org/datasets/add_dataset) for
    instructions.
*   Issues, feature requests, bugs,... have a much bigger impact than adding new
    datasets, as they benefit the entire TFDS community. See the
    [potential contribution list](https://github.com/tensorflow/datasets/issues?utf8=%E2%9C%93&q=is%3Aissue+is%3Aopen+-label%3A%22dataset+request%22+).
    Starts with the ones labeled with
    [contribution-welcome](https://github.com/tensorflow/datasets/issues?q=is%3Aissue+is%3Aopen+label%3A%22contributions+welcome%22)
    which are small self-contained easy issues to get started with.
*   Don't hesitate to take over bugs which are already assigned, but haven't
    been updated in a while.
*   No need to get the issue assigned to you. Simply comment on the issue when
    you're starting to work on it :)
*   Don't hesitate to ask for help if you're interested in an issue but don't
    know how to get started. And please send a draft PR if you want early
    feedback.
*   To avoid unnecessary duplication of work, check the list of
    [pending Pull Requests](https://github.com/tensorflow/datasets/pulls), and
    comment on issues you're working on.

## Setup

### Cloning the repo

To get started, clone or download the
[Tensorflow Datasets](https://github.com/tensorflow/datasets) repository and
install the repo locally.

```sh
git clone https://github.com/tensorflow/datasets.git
cd datasets/
```

Install the development dependencies:

```sh
pip install -e .  # Install minimal deps to use tensorflow_datasets
pip install -e ".[dev]"  # Install all deps required for testing and development
```

Note there is also a `pip install -e ".[tests-all]"` to install all
dataset-specific deps.

### Visual Studio Code

When developing with [Visual Studio Code](https://code.visualstudio.com/), our
repo comes with some
[pre-defined settings](https://github.com/tensorflow/datasets/tree/master/.vscode/settings.json)
to help development (correct indentation, pylint,...).

Note: enabling test discovery in VS Code may fails due to some VS Code bugs
[#13301](https://github.com/microsoft/vscode-python/issues/13301) and
[#6594](https://github.com/microsoft/vscode-python/issues/6594). To solve the
issues, you can look at the test discovery logs:

*   If you are encountering some TensorFlow warning message, try
    [this fix](https://github.com/microsoft/vscode-python/issues/6594#issuecomment-555680813).
*   If discovery fail due to missing import which should have been installed,
    please send a PR to update the `dev` pip install.

## PR checklist

### Sign the CLA

Contributions to this project must be accompanied by a Contributor License
Agreement (CLA). You (or your employer) retain the copyright to your
contribution; this simply gives us permission to use and redistribute your
contributions as part of the project. Head over to
<https://cla.developers.google.com/> to see your current agreements on file or
to sign a new one.

You generally only need to submit a CLA once, so if you've already submitted one
(even if it was for a different project), you probably don't need to do it
again.

### Follow best practices

*   Readability is important. Code should follow best programming practices
    (avoid duplication, factorise into small self-contained functions, explicit
    variables names,...)
*   Simpler is better (e.g. implementation split into multiple smaller
    self-contained PRs which is easier to review).
*   Add tests when required, existing tests should be passing.
*   Add [typing annotations](https://docs.python.org/3/library/typing.html)

### Check your style guide

Our style is based on
[Google Python Style Guide](https://github.com/google/styleguide/blob/gh-pages/pyguide.md),
which is based on
[PEP 8 Python style guide](https://www.python.org/dev/peps/pep-0008). New code
should try to follow
[Black code style](https://github.com/psf/black/blob/master/docs/the_black_code_style.md)
but with:

*   Line length: 80
*   2 spaces indentation instead of 4.
*   Single quote `'`

**Important:** Make sure to run `pylint` on your code to check your code is
properly formatted:

```sh
pip install pylint --upgrade
pylint tensorflow_datasets/core/some_file.py
```

You can try `yapf` to auto-format a file, but the tool is not perfect, so you'll
likely have to manually apply fixes afterward.

```sh
yapf tensorflow_datasets/core/some_file.py
```

Both `pylint` and `yapf` should have been installed with `pip install -e
".[dev]"` but can also be manually installed with `pip install`. If you're using
VS Code, those tools should be integrated in the UI.

### Docstrings and typing annotations

Classes and functions should be documented with docstrings and typing
annotation. Docstrings should follow the
[Google style](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods).
For example:

```python
def function(x: List[T]) -> T:
  """One line doc should end by a dot.

  * Use `backticks` for code and tripple backticks for multi-line.
  * Use full API name (`tfds.core.DatasetBuilder` instead of `DatasetBuilder`)
  * Use `Args:`, `Returns:`, `Yields:`, `Attributes:`, `Raises:`

  Args:
    x: description

  Returns:
    y: description
  """
```

### Add and run unittests

Make sure new features are tested with unit-tests. You can run tests through the
VS Code interface, or command line. For instance:

```sh
pytest -vv tensorflow_datasets/core/
```

`pytest` vs `unittest`: Historically, we have been using `unittest` module to
write tests. New tests should preferably use `pytest` which is more simple,
flexible, modern and used by most famous libraries (numpy, pandas, sklearn,
matplotlib, scipy, six,...). You can read the
[pytest guide](https://docs.pytest.org/en/stable/getting-started.html#getstarted)
if you're not familiar with pytest.

Tests for DatasetBuilders are special and are documented in the
[guide to add a dataset](https://github.com/tensorflow/datasets/blob/master/docs/add_dataset.md#test-your-dataset).

### Send the PR for reviews!

Congrats! See
[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
information on using pull requests.

---

### Datasets Versioning

# Datasets versioning

## Definition

Versioning can refer to different meaning:

*   The TFDS API version (pip version): `tfds.__version__`
*   The public dataset version, independent from TFDS (e.g.
    [Voc2007](https://pjreddie.com/projects/pascal-voc-dataset-mirror/),
    Voc2012). In TFDS each public dataset version should be implemented as an
    independent dataset:
    *   Either through
        [builder configs](https://www.tensorflow.org/datasets/add_dataset#dataset_configurationvariants_tfdscorebuilderconfig):
        E.g. `voc/2007`, `voc/2012`
    *   Either as 2 independent datasets: E.g. `wmt13_translate`,
        `wmt14_translate`
*   The dataset generation code version in TFDS (`my_dataset:1.0.0`): For
    example, if a bug is found in the TFDS implementation of `voc/2007`, the
    `voc.py` generation code will be updated (`voc/2007:1.0.0` ->
    `voc/2007:2.0.0`).

The rest of this guide only focus on the last definition (dataset code version
in the TFDS repository).

## Supported versions

As a general rule:

*   Only the last current version can be generated.
*   All previously generated dataset can be read (note: This require datasets
    generated with TFDS 4+).

```python
builder = tfds.builder('my_dataset')
builder.info.version  # Current version is: '2.0.0'

# download and load the last available version (2.0.0)
ds = tfds.load('my_dataset')

# Explicitly load a previous version (only works if
# `~/tensorflow_datasets/my_dataset/1.0.0/` already exists)
ds = tfds.load('my_dataset:1.0.0')
```

## Semantic

Every `DatasetBuilder` defined in TFDS comes with a version, for example:

```python
class MNIST(tfds.core.GeneratorBasedBuilder):
  VERSION = tfds.core.Version('2.0.0')
  RELEASE_NOTES = {
      '1.0.0': 'Initial release',
      '2.0.0': 'Update dead download url',
  }
```

The version follows
[Semantic Versioning 2.0.0](https://semver.org/spec/v2.0.0.html):
`MAJOR.MINOR.PATCH`. The purpose of the version is to be able to guarantee
reproducibility: loading a given dataset at a fixed version yields the same
data. More specifically:

 - If `PATCH` version is incremented, data as read by the client is the same,
 although data might be serialized differently on disk, or the metadata might
 have changed. For any given slice, the slicing API returns the same set of
 records.
 - If `MINOR` version is incremented, existing data as read by the client is the
 same, but there is additional data (features in each record). For any given
 slice, the  slicing API returns the same set of records.
 - If `MAJOR` version is incremented, the existing data has been changed and/or
 the slicing API doesn't necessarily return the same set of records for a given
 slice.

When a code change is made to the TFDS library and that code change impacts the
way a dataset is being serialized and/or read by the client, then the
corresponding builder version is incremented according to the above guidelines.

Note that the above semantic is best effort, and there might be un-noticed bugs
impacting a dataset while the version was not incremented. Such bugs are
eventually fixed, but if you heavily rely on the versioning, we advise you to
use TFDS from a released version (as opposed to `HEAD`).

Also note that some datasets have another versioning scheme independent from
the TFDS version. For example, the Open Images dataset has several versions,
and in TFDS, the corresponding builders are `open_images_v4`, `open_images_v5`,
...

## Loading a specific version

When loading a dataset or a `DatasetBuilder`, you can specify the version to
use. For example:

```python
tfds.load('imagenet2012:2.0.1')
tfds.builder('imagenet2012:2.0.1')

tfds.load('imagenet2012:2.0.0')  # Error: unsupported version.

# Resolves to 3.0.0 for now, but would resolve to 3.1.1 if when added.
tfds.load('imagenet2012:3.*.*')
```

If using TFDS for a publication, we advise you to:

 - **fix the `MAJOR` component of the version only**;
 - **advertise which version of the dataset was used in your results.**

Doing so should make it easier for your future self, your readers and
reviewers to reproduce your results.

## BUILDER_CONFIGS and versions

Some datasets define several `BUILDER_CONFIGS`. When that is the case, `version`
and `supported_versions` are defined on the config objects themselves. Other
than that, semantics and usage are identical. For example:

```python
class OpenImagesV4(tfds.core.GeneratorBasedBuilder):

  BUILDER_CONFIGS = [
      OpenImagesV4Config(
          name='original',
          version=tfds.core.Version('0.2.0'),
          supported_versions=[
            tfds.core.Version('1.0.0', "Major change in data"),
          ],
          description='Images at their original resolution and quality.'),
      ...
  ]

tfds.load('open_images_v4/original:1.*.*')
```

## Experimental version

Note: The following is bad practice, error prone and should be discouraged.

It is possible to allow 2 versions to be generated at the same time. One default
and one experimental version. For example:

```python
class MNIST(tfds.core.GeneratorBasedBuilder):
  VERSION = tfds.core.Version("1.0.0")  # Default version
  SUPPORTED_VERSIONS = [
      tfds.core.Version("2.0.0"),  # Experimental version
  ]


# Download and load default version 1.0.0
builder = tfds.builder('mnist')

#  Download and load experimental version 2.0.0
builder = tfds.builder('mnist', version='experimental_latest')
```

In the code, you need to make sure to support the 2 versions:

```python
class MNIST(tfds.core.GeneratorBasedBuilder):

  ...

  def _generate_examples(self, path):
    if self.info.version >= '2.0.0':
      ...
    else:
      ...
```

---

### Decode

# Customizing feature decoding

The `tfds.decode` API allows you override the default feature decoding. The main
use case is to skip the image decoding for better performance.

Note: This API gives you access to the low-level `tf.train.Example` format on
disk (as defined by the `FeatureConnector`). This API is targeted towards
advanced users who want better read performance with images.

## Usage examples

### Skipping the image decoding

To keep full control over the decoding pipeline, or to apply a filter before the
images get decoded (for better performance), you can skip the image decoding
entirely. This works with both `tfds.features.Image` and `tfds.features.Video`.

```python
ds = tfds.load('imagenet2012', split='train', decoders={
    'image': tfds.decode.SkipDecoding(),
})

for example in ds.take(1):
  assert example['image'].dtype == tf.string  # Images are not decoded
```

### Filter/shuffle dataset before images get decoded

Similarly to the previous example, you can use `tfds.decode.SkipDecoding()` to
insert additional `tf.data` pipeline customization before decoding the image.
That way the filtered images won't be decoded and you can use a bigger shuffle
buffer.

```python
# Load the base dataset without decoding
ds, ds_info = tfds.load(
    'imagenet2012',
    split='train',
    decoders={
        'image': tfds.decode.SkipDecoding(),  # Image won't be decoded here
    },
    as_supervised=True,
    with_info=True,
)
# Apply filter and shuffle
ds = ds.filter(lambda image, label: label != 10)
ds = ds.shuffle(10000)
# Then decode with ds_info.features['image']
ds = ds.map(
    lambda image, label: ds_info.features['image'].decode_example(image), label)

```

### Cropping and decoding at the same time

To override the default `tf.io.decode_image` operation, you can create a new
`tfds.decode.Decoder` object using the `tfds.decode.make_decoder()` decorator.

```python
@tfds.decode.make_decoder()
def decode_example(serialized_image, feature):
  crop_y, crop_x, crop_height, crop_width = 10, 10, 64, 64
  return tf.image.decode_and_crop_jpeg(
      serialized_image,
      [crop_y, crop_x, crop_height, crop_width],
      channels=feature.feature.shape[-1],
  )

ds = tfds.load('imagenet2012', split='train', decoders={
    # With video, decoders are applied to individual frames
    'image': decode_example(),
})
```

Which is equivalent to:

```python
def decode_example(serialized_image, feature):
  crop_y, crop_x, crop_height, crop_width = 10, 10, 64, 64
  return tf.image.decode_and_crop_jpeg(
      serialized_image,
      [crop_y, crop_x, crop_height, crop_width],
      channels=feature.shape[-1],
  )

ds, ds_info = tfds.load(
    'imagenet2012',
    split='train',
    with_info=True,
    decoders={
        'image': tfds.decode.SkipDecoding(),  # Skip frame decoding
    },
)
ds = ds.map(functools.partial(decode_example, feature=ds_info.features['image']))
```

### Customizing video decoding

Video are `Sequence(Image())`. When applying custom decoders, they will be
applied to individual frames. This mean decoders for images are automatically
compatible with video.

```python
@tfds.decode.make_decoder()
def decode_example(serialized_image, feature):
  crop_y, crop_x, crop_height, crop_width = 10, 10, 64, 64
  return tf.image.decode_and_crop_jpeg(
      serialized_image,
      [crop_y, crop_x, crop_height, crop_width],
      channels=feature.feature.shape[-1],
  )

ds = tfds.load('ucf101', split='train', decoders={
    # With video, decoders are applied to individual frames
    'video': decode_example(),
})
```

Which is equivalent to:

```python
def decode_frame(serialized_image):
  """Decodes a single frame."""
  crop_y, crop_x, crop_height, crop_width = 10, 10, 64, 64
  return tf.image.decode_and_crop_jpeg(
      serialized_image,
      [crop_y, crop_x, crop_height, crop_width],
      channels=ds_info.features['video'].shape[-1],
  )


def decode_video(example):
  """Decodes all individual frames of the video."""
  video = example['video']
  video = tf.map_fn(
      decode_frame,
      video,
      dtype=ds_info.features['video'].dtype,
      parallel_iterations=10,
  )
  example['video'] = video
  return example


ds, ds_info = tfds.load('ucf101', split='train', with_info=True, decoders={
    'video': tfds.decode.SkipDecoding(),  # Skip frame decoding
})
ds = ds.map(decode_video)  # Decode the video
```

### Only decode a sub-set of the features.

It's also possible to entirely skip some features by specifying only the
features you need. All other features will be ignored/skipped.

```python
builder = tfds.builder('my_dataset')
builder.as_dataset(split='train', decoders=tfds.decode.PartialDecoding({
    'image': True,
    'metadata': {'num_objects', 'scene_name'},
    'objects': {'label'},
})
```

TFDS will select the subset of `builder.info.features` matching the given
`tfds.decode.PartialDecoding` structure.

In the above code, the featured are implicitly extracted to match
`builder.info.features`. It is also possible to explicitly define the features.
The above code is equivalent to:

```python
builder = tfds.builder('my_dataset')
builder.as_dataset(split='train', decoders=tfds.decode.PartialDecoding({
    'image': tfds.features.Image(),
    'metadata': {
        'num_objects': tf.int64,
        'scene_name': tfds.features.Text(),
    },
    'objects': tfds.features.Sequence({
        'label': tfds.features.ClassLabel(names=[]),
    }),
})
```

The original metadata (label names, image shape,...) are automatically reused so
it's not required to provide them.

`tfds.decode.SkipDecoding` can be passed to `tfds.decode.PartialDecoding`,
through the `PartialDecoding(..., decoders={})` kwargs.

---

### External Tfrecord

# Load external tfrecord with TFDS


If you have a `tf.train.Example` proto (inside `.tfrecord`, `.riegeli`,...),
which has been generated by third party tools, that you would like to directly
load with tfds API, then this page is for you.

In order to load your `.tfrecord` files, you only need to:

*   Follow the TFDS naming convention.
*   Add metadata files (`dataset_info.json`, `features.json`) along your
    tfrecord files.

Limitations:

*   `tf.train.SequenceExample` is not supported, only `tf.train.Example`.
*   You need to be able to express the `tf.train.Example` in terms of
    `tfds.features` (see section below).

## File naming convention

TFDS supports defining a template for file names, which provides flexibility to
use different file naming schemes. The template is represented by a
`tfds.core.ShardedFileTemplate` and supports the following variables:
`{DATASET}`, `{SPLIT}`, `{FILEFORMAT}`, `{SHARD_INDEX}`, `{NUM_SHARDS}`, and
`{SHARD_X_OF_Y}`. For example, the default file naming scheme of TFDS is:
`{DATASET}-{SPLIT}.{FILEFORMAT}-{SHARD_X_OF_Y}`. For MNIST, this means that
[file names](https://console.cloud.google.com/storage/browser/tfds-data/datasets/mnist/3.0.1)
look as follows:

*   `mnist-test.tfrecord-00000-of-00001`
*   `mnist-train.tfrecord-00000-of-00001`


## Add metadata

### Provide the feature structure

For TFDS to be able to decode the `tf.train.Example` proto, you need to provide
the `tfds.features` structure matching your specs. For example:

```python
features = tfds.features.FeaturesDict({
    'image':
        tfds.features.Image(
            shape=(256, 256, 3),
            doc='Picture taken by smartphone, downscaled.'),
    'label':
        tfds.features.ClassLabel(names=['dog', 'cat']),
    'objects':
        tfds.features.Sequence({
            'camera/K': tfds.features.Tensor(shape=(3,), dtype=tf.float32),
        }),
})
```

Corresponds to the following `tf.train.Example` specs:

```python
{
    'image': tf.io.FixedLenFeature(shape=(), dtype=tf.string),
    'label': tf.io.FixedLenFeature(shape=(), dtype=tf.int64),
    'objects/camera/K': tf.io.FixedLenSequenceFeature(shape=(3,), dtype=tf.int64),
}
```

Specifying the features allow TFDS to automatically decode images, video,...
Like any other TFDS datasets, features metadata (e.g. label names,...) will be
exposed to the user (e.g. `info.features['label'].names`).

#### If you control the generation pipeline

If you generate datasets outside of TFDS but still control the generation
pipeline, you can use `tfds.features.FeatureConnector.serialize_example` to
encode your data from `dict[np.ndarray]` to `tf.train.Example` proto `bytes`:

```python
with tf.io.TFRecordWriter('path/to/file.tfrecord') as writer:
  for ex in all_exs:
    ex_bytes = features.serialize_example(data)
    writer.write(ex_bytes)
```

This will ensure feature compatibility with TFDS.

Similarly, a `feature.deserialize_example` exists to decode the proto
([example](https://www.tensorflow.org/datasets/features#serializedeserialize_to_proto))

#### If you don't control the generation pipeline

If you want to see how `tfds.features` are represented in a `tf.train.Example`,
you can examine this in colab:

*   To translate `tfds.features` into the human readable structure of the
    `tf.train.Example`, you can call `features.get_serialized_info()`.
*   To get the exact `FixedLenFeature`,... spec passed to
    `tf.io.parse_single_example`, you can use `spec = features.tf_example_spec`

Note: If you're using custom feature connector, make sure to implement
`to_json_content`/`from_json_content` and test with `self.assertFeature` (see
[feature connector guide](https://www.tensorflow.org/datasets/features#create_your_own_tfdsfeaturesfeatureconnector))

### Get statistics on splits


TFDS requires to know the exact number of examples within each shard. This is
required for features like `len(ds)`, or the
[subplit API](https://www.tensorflow.org/datasets/splits):
`split='train[75%:]'`.

*   If you have this information, you can explicitly create a list of
    `tfds.core.SplitInfo` and skip to the next section:

    ```python
    split_infos = [
        tfds.core.SplitInfo(
            name='train',
            shard_lengths=[1024, ...],  # Num of examples in shard0, shard1,...
            num_bytes=0,  # Total size of your dataset (if unknown, set to 0)
        ),
        tfds.core.SplitInfo(name='test', ...),
    ]
    ```

*   If you do not know this information, you can compute it using the
    `compute_split_info.py` script (or in your own script with
    `tfds.folder_dataset.compute_split_info`). It will launch a beam pipeline
    which will read all shards on the given directory and compute the info.


### Add metadata files

To automatically add the proper metadata files along your dataset, use
`tfds.folder_dataset.write_metadata`:

```python
tfds.folder_dataset.write_metadata(
    data_dir='/path/to/my/dataset/1.0.0/',
    features=features,
    # Pass the `out_dir` argument of compute_split_info (see section above)
    # You can also explicitly pass a list of `tfds.core.SplitInfo`.
    split_infos='/path/to/my/dataset/1.0.0/',
    # Pass a custom file name template or use None for the default TFDS
    # file name template.
    filename_template='{SPLIT}-{SHARD_X_OF_Y}.{FILEFORMAT}',

    # Optionally, additional DatasetInfo metadata can be provided
    # See:
    # https://www.tensorflow.org/datasets/api_docs/python/tfds/core/DatasetInfo
    description="""Multi-line description."""
    homepage='http://my-project.org',
    supervised_keys=('image', 'label'),
    citation="""BibTex citation.""",
)
```

Once the function has been called once on your dataset directory, metadata files
( `dataset_info.json`,...) have been added and your datasets are ready to be
loaded with TFDS (see next section).

## Load dataset with TFDS

### Directly from folder

Once the metadata have been generated, datasets can be loaded using
`tfds.builder_from_directory` which returns a `tfds.core.DatasetBuilder` with
the standard TFDS API (like `tfds.builder`):

```python
builder = tfds.builder_from_directory('~/path/to/my_dataset/3.0.0/')

# Metadata are available as usual
builder.info.splits['train'].num_examples

# Construct the tf.data.Dataset pipeline
ds = builder.as_dataset(split='train[75%:]')
for ex in ds:
  ...
```

### Directly from multiple folders

It is also possible to load data from multiple folders. This can happen, for
example, in reinforcement learning when multiple agents are each generating a
separate dataset and you want to load all of them together. Other use cases are
when a new dataset is produced on a regular basis, e.g. a new dataset per day,
and you want to load data from a date range.

To load data from multiple folders, use `tfds.builder_from_directories`, which
returns a `tfds.core.DatasetBuilder` with the standard TFDS API (like
`tfds.builder`):

```python
builder = tfds.builder_from_directories(builder_dirs=[
    '~/path/my_dataset/agent1/1.0.0/',
    '~/path/my_dataset/agent2/1.0.0/',
    '~/path/my_dataset/agent3/1.0.0/',
])

# Metadata are available as usual
builder.info.splits['train'].num_examples

# Construct the tf.data.Dataset pipeline
ds = builder.as_dataset(split='train[75%:]')
for ex in ds:
  ...
```

Note: each folder must have its own metadata, because this contains information
about the splits.

### Folder structure (optional)

For better compatibility with TFDS, you can organize your data as
`<data_dir>/<dataset_name>[/<dataset_config>]/<dataset_version>`. For example:

```
data_dir/
    dataset0/
        1.0.0/
        1.0.1/
    dataset1/
        config0/
            2.0.0/
        config1/
            2.0.0/
```

This will make your datasets compatible with the `tfds.load` / `tfds.builder`
API, simply by providing `data_dir/`:

```python
ds0 = tfds.load('dataset0', data_dir='data_dir/')
ds1 = tfds.load('dataset1/config0', data_dir='data_dir/')
```

---

### Features

# FeatureConnector

The `tfds.features.FeatureConnector` API:

*   Defines the structure, shapes, dtypes of the final `tf.data.Dataset`
*   Abstract away serialization to/from disk.
*   Expose additional metadata (e.g. label names, audio sample rate,...)

## Overview

`tfds.features.FeatureConnector` defines the dataset features structure (in
`tfds.core.DatasetInfo`):

```python
tfds.core.DatasetInfo(
    features=tfds.features.FeaturesDict({
        'image': tfds.features.Image(shape=(28, 28, 1), doc='Grayscale image'),
        'label': tfds.features.ClassLabel(
            names=['no', 'yes'],
            doc=tfds.features.Documentation(
                desc='Whether this is a picture of a cat',
                value_range='yes or no'
            ),
        ),
        'metadata': {
            'id': tf.int64,
            'timestamp': tfds.features.Scalar(
                tf.int64,
                doc='Timestamp when this picture was taken as seconds since epoch'),
            'language': tf.string,
        },
    }),
)
```

Features can be documented by either using just a textual description
(`doc='description'`) or by using `tfds.features.Documentation` directly to
provide a more detailed feature description.

Features can be:

*   Scalar values: `tf.bool`, `tf.string`, `tf.float32`,... When you want to
    document the feature, you can also use `tfds.features.Scalar(tf.int64,
    doc='description')`.
*   `tfds.features.Audio`, `tfds.features.Video`,... (see
    [the list](https://www.tensorflow.org/datasets/api_docs/python/tfds/features?version=nightly)
    of available features)
*   Nested `dict` of features: `{'metadata': {'image': Image(), 'description':
    tf.string}}`,...
*   Nested `tfds.features.Sequence`: `Sequence({'image': ..., 'id': ...})`,
    `Sequence(Sequence(tf.int64))`,...

During generation, the examples will be automatically serialized by
`FeatureConnector.encode_example` into a format suitable to disk (currently
`tf.train.Example` protocol buffers):

```python
yield {
    'image': '/path/to/img0.png',  # `np.array`, file bytes,... also accepted
    'label': 'yes',  # int (0-num_classes) also accepted
    'metadata': {
        'id': 43,
        'language': 'en',
    },
}
```

When reading the dataset (e.g. with `tfds.load`), the data is automtically
decoded with `FeatureConnector.decode_example`. The returned `tf.data.Dataset`
will match the `dict` structure defined in `tfds.core.DatasetInfo`:

```python
ds = tfds.load(...)
ds.element_spec == {
    'image': tf.TensorSpec(shape=(28, 28, 1), tf.uint8),
    'label': tf.TensorSpec(shape=(), tf.int64),
    'metadata': {
        'id': tf.TensorSpec(shape=(), tf.int64),
        'language': tf.TensorSpec(shape=(), tf.string),
    },
}
```

## Serialize/deserialize to proto

TFDS expose a low-level API to serialize/deserialize examples to
`tf.train.Example` proto.

To serialize `dict[np.ndarray | Path | str | ...]` to proto `bytes`, use
`features.serialize_example`:

```python
with tf.io.TFRecordWriter('path/to/file.tfrecord') as writer:
  for ex in all_exs:
    ex_bytes = features.serialize_example(data)
    f.write(ex_bytes)
```

To deserialize to proto `bytes` to `tf.Tensor`, use
`features.deserialize_example`:

```python
ds = tf.data.TFRecordDataset('path/to/file.tfrecord')
ds = ds.map(features.deserialize_example)
```

## Access metadata

See the
[introduction doc](https://www.tensorflow.org/datasets/overview#access_the_dataset_metadata)
to access features metadata (label names, shape, dtype,...). Example:

```python
ds, info = tfds.load(..., with_info=True)

info.features['label'].names  # ['cat', 'dog', ...]
info.features['label'].str2int('cat')  # 0
```

## Create your own `tfds.features.FeatureConnector`

If you believe a feature is missing from the
[available features](https://www.tensorflow.org/datasets/api_docs/python/tfds/features#classes),
please open a [new issue](https://github.com/tensorflow/datasets/issues).

To create your own feature connector, you need to inherit from
`tfds.features.FeatureConnector` and implement the abstract methods.

*   If your feature is a single tensor value, it's best to inherit from
    `tfds.features.Tensor` and use `super()` when needed. See
    `tfds.features.BBoxFeature` source code for an example.
*   If your feature is a container of multiple tensors, it's best to inherit
    from `tfds.features.FeaturesDict` and use the `super()` to automatically
    encode sub-connectors.

The `tfds.features.FeatureConnector` object abstracts away how the feature is
encoded on disk from how it is presented to the user. Below is a diagram showing
the abstraction layers of the dataset and the transformation from the raw
dataset files to the `tf.data.Dataset` object.

<p align="center">
  <img src="dataset_layers.png" alt="DatasetBuilder abstraction layers" width="700"/>
</p>

To create your own feature connector, subclass `tfds.features.FeatureConnector`
and implement the abstract methods:

*   `encode_example(data)`: Defines how to encode the data given in the
    generator `_generate_examples()` into a `tf.train.Example` compatible data.
    Can return a single value, or a `dict` of values.
*   `decode_example(data)`: Defines how to decode the data from the tensor read
    from `tf.train.Example` into user tensor returned by `tf.data.Dataset`.
*   `get_tensor_info()`: Indicates the shape/dtype of the tensor(s) returned by
    `tf.data.Dataset`. May be optional if inheriting from another
    `tfds.features`.
*   (optionally) `get_serialized_info()`: If the info returned by
    `get_tensor_info()` is different from how the data are actually written on
    disk, then you need to overwrite `get_serialized_info()` to match the specs
    of the `tf.train.Example`
*   `to_json_content`/`from_json_content`: This is required to allow your
    dataset to be loaded without the original source code. See
    [Audio feature](https://github.com/tensorflow/datasets/blob/65a76cb53c8ff7f327a3749175bc4f8c12ff465e/tensorflow_datasets/core/features/audio_feature.py#L121)
    for an example.

Note: Make sure to test your Feature connectors with `self.assertFeature` and
`tfds.testing.FeatureExpectationItem`. Have a look at
[test examples](https://github.com/tensorflow/datasets/tree/master/tensorflow_datasets/core/features/image_feature_test.py):

For more info, have a look at `tfds.features.FeatureConnector` documentation.
It's also best to look at
[real examples](https://github.com/tensorflow/datasets/tree/master/tensorflow_datasets/core/features).

---

### Format Specific Dataset Builders

# Format-specific Dataset Builders

[TOC]

This guide documents all format-specific dataset builders currently available in
TFDS.

Format-specific dataset builders are subclasses of
[`tfds.core.GeneratorBasedBuilder`](https://www.tensorflow.org/datasets/api_docs/python/tfds/core/GeneratorBasedBuilder)
which take care of most data processing for a specific data format.

## Datasets based on `tf.data.Dataset`

If you want to create a TFDS dataset from a dataset that's in `tf.data.Dataset`
([reference](https://www.tensorflow.org/api_docs/python/tf/data/Dataset))
format, then you can use `tfds.dataset_builders.TfDataBuilder` (see
[API docs](https://www.tensorflow.org/datasets/api_docs/python/tfds/dataset_builders/TfDataBuilder)).

We envision two typical uses of this class:

*   Creating experimental datasets in a notebook-like environment
*   Defining a dataset builder in code

### Creating a new dataset from a notebook

Suppose you are working in a notebook, loaded some data as a `tf.data.Dataset`,
applied various transformations (map, filter, etc) and now you want to store
this data and easily share it with teammates or load it in other notebooks.
Instead of having to define a new dataset builder class, you can also
instantiate a `tfds.dataset_builders.TfDataBuilder` and call
`download_and_prepare` to store your dataset as a TFDS dataset.

Because it's a TFDS dataset, you can version it, use configs, have different
splits, and document it for easier use later. This means that you also have to
tell TFDS what the features are in your dataset.

Here's a dummy example of how you can use it.

```python
import tensorflow as tf
import tensorflow_datasets as tfds

my_ds_train = tf.data.Dataset.from_tensor_slices({"number": [1, 2, 3]})
my_ds_test = tf.data.Dataset.from_tensor_slices({"number": [4, 5]})

# Optionally define a custom `data_dir`.
# If None, then the default data dir is used.
custom_data_dir = "/my/folder"

# Define the builder.
single_number_builder = tfds.dataset_builders.TfDataBuilder(
    name="my_dataset",
    config="single_number",
    version="1.0.0",
    data_dir=custom_data_dir,
    split_datasets={
        "train": my_ds_train,
        "test": my_ds_test,
    },
    features=tfds.features.FeaturesDict({
        "number": tfds.features.Scalar(dtype=tf.int64),
    }),
    description="My dataset with a single number.",
    release_notes={
        "1.0.0": "Initial release with numbers up to 5!",
    }
)

# Make the builder store the data as a TFDS dataset.
single_number_builder.download_and_prepare()
```

The `download_and_prepare` method will iterate over the input `tf.data.Dataset`s
and store the corresponding TFDS dataset in
`/my/folder/my_dataset/single_number/1.0.0`, which will contain both the train
and test splits.

The `config` argument is optional and can be useful if you want to store
different configs under the same dataset.

The `data_dir` argument can be used to store the generated TFDS dataset in a
different folder, for example in your own sandbox if you don't want to share
this with others (yet). Note that when doing this, you also need to pass the
`data_dir` to `tfds.load`. If the `data_dir` argument is not specified, then the
default TFDS data dir will be used.

#### Loading your dataset

After the TFDS dataset has been stored, it can be loaded from other scripts or
by teammates if they have access to the data:

```python
# If no custom data dir was specified:
ds_test = tfds.load("my_dataset/single_number", split="test")

# When there are multiple versions, you can also specify the version.
ds_test = tfds.load("my_dataset/single_number:1.0.0", split="test")

# If the TFDS was stored in a custom folder, then it can be loaded as follows:
custom_data_dir = "/my/folder"
ds_test = tfds.load("my_dataset/single_number:1.0.0", split="test", data_dir=custom_data_dir)
```

#### Adding a new version or config

After iterating further on your dataset, you may have added or changed some of
the transformations of the source data. To store and share this dataset, you can
easily store this as a new version.

```python
def add_one(example):
  example["number"] = example["number"] + 1
  return example

my_ds_train_v2 = my_ds_train.map(add_one)
my_ds_test_v2 = my_ds_test.map(add_one)

single_number_builder_v2 = tfds.dataset_builders.TfDataBuilder(
    name="my_dataset",
    config="single_number",
    version="1.1.0",
    data_dir=custom_data_dir,
    split_datasets={
        "train": my_ds_train_v2,
        "test": my_ds_test_v2,
    },
    features=tfds.features.FeaturesDict({
        "number": tfds.features.Scalar(dtype=tf.int64, doc="Some number"),
    }),
    description="My dataset with a single number.",
    release_notes={
        "1.1.0": "Initial release with numbers up to 6!",
        "1.0.0": "Initial release with numbers up to 5!",
    }
)

# Make the builder store the data as a TFDS dataset.
single_number_builder_v2.download_and_prepare()
```

### Defining a new dataset builder class

You can also define a new `DatasetBuilder` based on this class.

```python
import tensorflow as tf
import tensorflow_datasets as tfds

class MyDatasetBuilder(tfds.dataset_builders.TfDataBuilder):
  def __init__(self):
    ds_train = tf.data.Dataset.from_tensor_slices([1, 2, 3])
    ds_test = tf.data.Dataset.from_tensor_slices([4, 5])
    super().__init__(
        name="my_dataset",
        version="1.0.0",
        split_datasets={
            "train": ds_train,
            "test": ds_test,
        },
        features=tfds.features.FeaturesDict({
            "number": tfds.features.Scalar(dtype=tf.int64),
        }),
        config="single_number",
        description="My dataset with a single number.",
        release_notes={
            "1.0.0": "Initial release with numbers up to 5!",
        })
```

## CroissantBuilder

### The format

[Croissant](https://github.com/mlcommons/croissant) 🥐 is a high-level format for
machine learning datasets that combines metadata, resource file descriptions,
data structure, and default ML semantics into a single file; it works with
existing datasets to make them easier to find, use, and support with tools.

Croissant builds on schema.org and its `sc:Dataset` vocabulary, a widely used
format to represent datasets on the Web, and make them searchable.

### `CroissantBuilder`

A `CroissantBuilder` defines a TFDS dataset based on a Croissant 🥐 metadata
file; each of the `record_set_ids` specified will result in a separate
`ConfigBuilder`.

For example, to initialize a `CroissantBuilder` for the MNIST dataset using its
[Croissant 🥐 definition](https://github.com/mlcommons/croissant/tree/main/datasets/1.0/huggingface-mnist):

```python
import tensorflow_datasets as tfds
builder = tfds.dataset_builders.CroissantBuilder(
    jsonld="https://raw.githubusercontent.com/mlcommons/croissant/main/datasets/0.8/huggingface-mnist/metadata.json",
    file_format='array_record',
)
builder.download_and_prepare()
ds = builder.as_data_source()
print(ds['default'][0])
```

## CoNLL

### The format

[CoNLL](https://aclanthology.org/W03-0419.pdf) is a popular format used to
represent annotated text data.

CoNLL-formatted data usually contain one token with its linguistic annotations
per line; within the same line, annotations are usually separated by spaces or
tabs. Empty lines represent sentence boundaries.

Consider as an example the following sentence from the
[conll2003](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/text/conll2003/conll2003.py)
dataset, which follows the CoNLL annotation format:

```markdown
U.N. NNP I-NP I-ORG official
NN I-NP O
Ekeus NNP I-NP I-PER
heads VBZ I-VP O
for IN I-PP O
Baghdad NNP I-NP
I-LOC . . O O
```

### `ConllDatasetBuilder`

To add a new CoNLL-based dataset to TFDS, you can base your dataset builder
class on `tfds.dataset_builders.ConllDatasetBuilder`. This base class contains
common code to deal with the specificities of CoNLL datasets (iterating over the
column-based format, precompiled lists of features and tags, ...).

`tfds.dataset_builders.ConllDatasetBuilder` implements a CoNLL-specific
`GeneratorBasedBuilder`. Refer to the following class as a minimal example of a
CoNLL dataset builder:

```python
from tensorflow_datasets.core.dataset_builders.conll import conll_dataset_builder_utils as conll_lib
import tensorflow_datasets.public_api as tfds

class MyCoNNLDataset(tfds.dataset_builders.ConllDatasetBuilder):
  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {'1.0.0': 'Initial release.'}

  # conllu_lib contains a set of ready-to-use CONLL-specific configs.
  BUILDER_CONFIGS = [conll_lib.CONLL_2003_CONFIG]

  def _info(self) -> tfds.core.DatasetInfo:
    return self.create_dataset_info(
        # ...
    )

  def _split_generators(self, dl_manager):
    path = dl_manager.download_and_extract('https://data-url')

    return {'train': self._generate_examples(path=path / 'train.txt'),
            'test': self._generate_examples(path=path / 'train.txt'),
    }
```

As for standard dataset builders, it requires to overwrite the class methods
`_info` and `_split_generators`. Depending on the dataset, you might need to
update also
[conll_dataset_builder_utils.py](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/core/dataset_builders/conll/conll_dataset_builder_utils.py)
to include the features and the list of tags specific to your dataset.

The `_generate_examples` method should not require further overwriting, unless
your dataset needs specific implementation.

### Examples

Consider
[conll2003](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/text/conll2003/conll2003.py)
as an example of a dataset implemented using the CoNLL-specific dataset builder.

### CLI

The easiest way to write a new CoNLL-based dataset is to use the
[TFDS CLI](https://www.tensorflow.org/datasets/cli):

```sh
cd path/to/my/project/datasets/
tfds new my_dataset --format=conll   # Create `my_dataset/my_dataset.py` CoNLL-specific template files
```

## CoNLL-U

### The format

[CoNLL-U](https://universaldependencies.org/format.html) is a popular format
used to represent annotated text data.

CoNLL-U enhances the CoNLL format by adding a number of features, such as
support for
[multi-token words](https://universaldependencies.org/u/overview/tokenization.html).
CoNLL-U formatted data usually contain one token with its linguistic annotations
per line; within the same line, annotations are usually separated by single tab
characters. Empty lines represent sentence boundaries.

Typically, each CoNLL-U annotated word line contains the following fields, as
reported in the
[official documentation](https://universaldependencies.org/format.html):

*   ID: Word index, integer starting at 1 for each new sentence; may be a range
    for multiword tokens; may be a decimal number for empty nodes (decimal
    numbers can be lower than 1 but must be greater than 0).
*   FORM: Word form or punctuation symbol.
*   LEMMA: Lemma or stem of word form.
*   UPOS: Universal part-of-speech tag.
*   XPOS: Language-specific part-of-speech tag; underscore if not available.
*   FEATS: List of morphological features from the universal feature inventory
    or from a defined language-specific extension; underscore if not available.
*   HEAD: Head of the current word, which is either a value of ID or zero (0).
*   DEPREL: Universal dependency relation to the HEAD (root iff HEAD = 0) or a
    defined language-specific subtype of one.
*   DEPS: Enhanced dependency graph in the form of a list of head-deprel pairs.
*   MISC: Any other annotation.

Consider as an example the following CoNLL-U annotated sentence from the
[official documentation](https://universaldependencies.org/format.html):

```markdown
1-2    vámonos   _
1      vamos     ir
2      nos       nosotros
3-4    al        _
3      a         a
4      el        el
5      mar       mar
```

### `ConllUDatasetBuilder`

To add a new CoNLL-U based dataset to TFDS, you can base your dataset builder
class on `tfds.dataset_builders.ConllUDatasetBuilder`. This base class contains
common code to deal with the specificities of CoNLL-U datasets (iterating over
the column-based format, precompiled lists of features and tags, ...).

`tfds.dataset_builders.ConllUDatasetBuilder` implements a CoNLL-U specific
`GeneratorBasedBuilder`. Refer to the following class as a minimal example of a
CoNLL-U dataset builder:

```python
from tensorflow_datasets.core.dataset_builders.conll import conllu_dataset_builder_utils as conllu_lib
import tensorflow_datasets.public_api as tfds

class MyCoNNLUDataset(tfds.dataset_builders.ConllUDatasetBuilder):
  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {'1.0.0': 'Initial release.'}

  # conllu_lib contains a set of ready-to-use features.
  BUILDER_CONFIGS = [
      conllu_lib.get_universal_morphology_config(
          language='en',
          features=conllu_lib.UNIVERSAL_DEPENDENCIES_FEATURES,
      )
  ]

  def _info(self) -> tfds.core.DatasetInfo:
    return self.create_dataset_info(
        # ...
    )

  def _split_generators(self, dl_manager):
    path = dl_manager.download_and_extract('https://data-url')

    return {
        'train':
            self._generate_examples(
                path=path / 'train.txt',
                # If necessary, add optional custom processing (see conllu_lib
                # for examples).
                # process_example_fn=...,
            )
    }
```

As for standard dataset builders, it requires to overwrite the class methods
`_info` and `_split_generators`. Depending on the dataset, you might need to
update also
[conllu_dataset_builder_utils.py](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/core/dataset_builders/conll/conllu_dataset_builder_utils.py)
to include the features and the list of tags specific to your dataset.

The `_generate_examples` method should not require further overwriting, unless
your dataset needs specific implementation. Note that, if your dataset requires
specific preprocessing - for example if it considers non-classic
[universal dependency features](https://universaldependencies.org/guidelines.html) -
you might need to update the `process_example_fn` attribute of your
[`generate_examples`](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/core/dataset_builders/conll/conllu_dataset_builder.py#L192)
function (see the
[xtreme_pos](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/text/universal_dependencies/xtreme_pos.py)
daset as an example).

### Examples

Consider the following datasets, which use the CoNNL-U specific dataset builder,
as examples:

*   [universal_dependencies](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/text/universal_dependencies/universal_dependencies.py)
*   [xtreme_pos](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/text/universal_dependencies/xtreme_pos.py)

### CLI

The easiest way to write a new CoNLL-U based dataset is to use the
[TFDS CLI](https://www.tensorflow.org/datasets/cli):

```sh
cd path/to/my/project/datasets/
tfds new my_dataset --format=conllu   # Create `my_dataset/my_dataset.py` CoNLL-U specific template files
```

---

### Gcs

# tfds and Google Cloud Storage

Google Cloud Storage (GCS) can be used with tfds for multiple reasons:

*   Storing preprocessed data
*   Accessing datasets that have data stored on GCS

## Access through TFDS GCS bucket

Some datasets are available directly in our GCS bucket
[`gs://tfds-data/datasets/`](https://console.cloud.google.com/storage/browser/tfds-data)
without any authentication:

*   If `tfds.load(..., try_gcs=False)` (default), the dataset will be copied
    locally in `~/tensorflow_datasets` during `download_and_prepare`.
*   If `tfds.load(..., try_gcs=True)`, the dataset will be streamed directly
    from GCS (`download_and_prepare` will be skipped).

You can check whether a dataset is hosted on the public bucket with
`tfds.is_dataset_on_gcs('mnist')`.

## Authentication

Before starting, you should decide on how you want to authenticate. There are
three options:

*   no authentication (a.k.a anonymous access)
*   using your Google account
*   using a service account (can be easily shared with others in your team)

You can find detailed information in
[Google Cloud documentation](https://cloud.google.com/docs/authentication/getting-started)

### Simplified instructions

If you run from colab, you can authenticate with your account, but running:

```python
from google.colab import auth
auth.authenticate_user()
```

If you run on your local machine (or in VM), you can authenticate with your
account by running:

```shell
gcloud auth application-default login
```

If you want to login with service account, download the JSON file key and set

```shell
export GOOGLE_APPLICATION_CREDENTIALS=<JSON_FILE_PATH>
```

## Using Google Cloud Storage to store preprocessed data

Normally when you use TensorFlow Datasets, the downloaded and prepared data will
be cached in a local directory (by default `~/tensorflow_datasets`).

In some environments where local disk may be ephemeral (a temporary cloud server
or a [Colab notebook](https://colab.research.google.com)) or you need the data
to be accessible by multiple machines, it's useful to set `data_dir` to a cloud
storage system, like a Google Cloud Storage (GCS) bucket.

### How?

[Create a GCS bucket](https://cloud.google.com/storage/docs/creating-buckets)
and ensure you (or your service account) have read/write permissions on it (see
authorization instructions above)

When you use `tfds`, you can set `data_dir` to `"gs://YOUR_BUCKET_NAME"`

```python
ds_train, ds_test = tfds.load(name="mnist", split=["train", "test"], data_dir="gs://YOUR_BUCKET_NAME")
```

### Caveats:

*   This approach works for datasets that only use `tf.io.gfile` for data
    access. This is true for most datasets, but not all.
*   Remember that accessing GCS is accessing a remote server and streaming data
    from it, so you may incur network costs.

## Accessing datasets stored on GCS

If dataset owners allowed anonymous access, you can just go ahead and run the
tfds.load code - and it would work like a normal internet download.

If dataset requires authentication, please use the instructions above to decide
on which option you want (own account vs service account) and communicate the
account name (a.k.a email) to the dataset owner. After they enable you access to
the GCS directory, you should be able to run the tfds download code.

---

### Performances

# Performance tips

This document provides TensorFlow Datasets (TFDS)-specific performance tips.
Note that TFDS provides datasets as `tf.data.Dataset` objects, so the advice
from the
[`tf.data` guide](https://www.tensorflow.org/guide/data_performance#optimize_performance)
still applies.

## Benchmark datasets

Use `tfds.benchmark(ds)` to benchmark any `tf.data.Dataset` object.

Make sure to indicate the `batch_size=` to normalize the results (e.g. 100
iter/sec -> 3200 ex/sec). This works with any iterable (e.g.
`tfds.benchmark(tfds.as_numpy(ds))`).

```python
ds = tfds.load('mnist', split='train').batch(32).prefetch()
# Display some benchmark statistics
tfds.benchmark(ds, batch_size=32)
# Second iteration is much faster, due to auto-caching
tfds.benchmark(ds, batch_size=32)
```

## Small datasets (less than 1 GB)

All TFDS datasets store the data on disk in the
[`TFRecord`](https://www.tensorflow.org/tutorials/load_data/tfrecord) format.
For small datasets (e.g. MNIST, CIFAR-10/-100), reading from `.tfrecord` can add
significant overhead.

As those datasets fit in memory, it is possible to significantly improve the
performance by caching or pre-loading the dataset. Note that TFDS automatically
caches small datasets (the following section has the details).

### Caching the dataset

Here is an example of a data pipeline which explicitly caches the dataset after
normalizing the images.

```python
def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label


ds, ds_info = tfds.load(
    'mnist',
    split='train',
    as_supervised=True,  # returns `(img, label)` instead of dict(image=, ...)
    with_info=True,
)
# Applying normalization before `ds.cache()` to re-use it.
# Note: Random transformations (e.g. images augmentations) should be applied
# after both `ds.cache()` (to avoid caching randomness) and `ds.batch()` (for
# vectorization [1]).
ds = ds.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.cache()
# For true randomness, we set the shuffle buffer to the full dataset size.
ds = ds.shuffle(ds_info.splits['train'].num_examples)
# Batch after shuffling to get unique batches at each epoch.
ds = ds.batch(128)
ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
```

*   [[1] Vectorizing mapping](https://www.tensorflow.org/guide/data_performance#vectorizing_mapping)

When iterating over this dataset, the second iteration will be much faster than
the first one thanks to the caching.

### Auto-caching

By default, TFDS auto-caches (with `ds.cache()`) datasets which satisfy the
following constraints:

*   Total dataset size (all splits) is defined and < 250 MiB
*   `shuffle_files` is disabled, or only a single shard is read

It is possible to opt out of auto-caching by passing `try_autocaching=False` to
`tfds.ReadConfig` in `tfds.load`. Have a look at the dataset catalog
documentation to see if a specific dataset will use auto-cache.

### Loading the full data as a single Tensor

If your dataset fits into memory, you can also load the full dataset as a single
Tensor or NumPy array. It is possible to do so by setting `batch_size=-1` to
batch all examples in a single `tf.Tensor`. Then use `tfds.as_numpy` for the
conversion from `tf.Tensor` to `np.array`.

```python
(img_train, label_train), (img_test, label_test) = tfds.as_numpy(tfds.load(
    'mnist',
    split=['train', 'test'],
    batch_size=-1,
    as_supervised=True,
))
```

## Large datasets

Large datasets are sharded (split in multiple files) and typically do not fit
in memory, so they should not be cached.

### Shuffle and training

During training, it's important to shuffle the data well - poorly shuffled data
can result in lower training accuracy.

In addition to using `ds.shuffle` to shuffle records, you should also set
`shuffle_files=True` to get good shuffling behavior for larger datasets that are
sharded into multiple files. Otherwise, epochs will read the shards in the same
order, and so data won't be truly randomized.

```python
ds = tfds.load('imagenet2012', split='train', shuffle_files=True)
```

Additionally, when `shuffle_files=True`, TFDS disables
[`options.deterministic`](https://www.tensorflow.org/api_docs/python/tf/data/Options#deterministic),
which may give a slight performance boost. To get deterministic shuffling, it is
possible to opt-out of this feature with `tfds.ReadConfig`: either by setting
`read_config.shuffle_seed` or overwriting `read_config.options.deterministic`.

### Auto-shard your data across workers (TF)

When training on multiple workers, you can use the `input_context` argument of
`tfds.ReadConfig`, so each worker will read a subset of the data.

```python
input_context = tf.distribute.InputContext(
    input_pipeline_id=1,  # Worker id
    num_input_pipelines=4,  # Total number of workers
)
read_config = tfds.ReadConfig(
    input_context=input_context,
)
ds = tfds.load('dataset', split='train', read_config=read_config)
```

This is complementary to the subsplit API. First, the subplit API is applied:
`train[:50%]` is converted into a list of files to read. Then, a `ds.shard()` op
is applied on those files. For example, when using `train[:50%]` with
`num_input_pipelines=2`, each of the 2 workers will read 1/4 of the data.

When `shuffle_files=True`, files are shuffled within one worker, but not across
workers. Each worker will read the same subset of files between epochs.

Note: When using `tf.distribute.Strategy`, the `input_context` can be
automatically created with
[distribute_datasets_from_function](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#distribute_datasets_from_function)

### Auto-shard your data across workers (Jax)

With Jax, you can use the `tfds.split_for_jax_process` or `tfds.even_splits` API
to distribute your data across workers. See the
[split API guide](https://www.tensorflow.org/datasets/splits).

```python
split = tfds.split_for_jax_process('train', drop_remainder=True)
ds = tfds.load('my_dataset', split=split)
```

`tfds.split_for_jax_process` is a simple alias for:

```python
# The current `process_index` loads only `1 / process_count` of the data.
splits = tfds.even_splits('train', n=jax.process_count(), drop_remainder=True)
split = splits[jax.process_index()]
```

### Faster image decoding

By default, TFDS automatically decodes images. However, there are cases where it
can be more performant to skip the image decoding with
`tfds.decode.SkipDecoding` and manually apply the `tf.io.decode_image` op:

*   When filtering examples (with `tf.data.Dataset.filter`), to decode images
    after examples have been filtered.
*   When cropping images, to use the fused `tf.image.decode_and_crop_jpeg` op.

The code for both examples is available in the
[decode guide](https://www.tensorflow.org/datasets/decode#usage_examples).

### Skip unused features

If you're only using a subset of the features, it is possible to entirely skip
some features. If your dataset has many unused features, not decoding them can
significantly improve performances. See
https://www.tensorflow.org/datasets/decode#only_decode_a_sub-set_of_the_features.

## tf.data uses all my RAM!

If you are limited in RAM, or if you are loading many datasets in parallel while
using `tf.data`, here are a few options which can help:

### Override buffer size

```py
builder.as_dataset(
  read_config=tfds.ReadConfig(
    ...
    override_buffer_size=1024,  # Save quite a bit of RAM.
  ),
  ...
)
```

This overrides the `buffer_size` passed to `TFRecordDataset` (or equivalent):
https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset#args.

### Use tf.data.Dataset.with_options to stop magic behaviors

https://www.tensorflow.org/api_docs/python/tf/data/Dataset#with_options

```py
options = tf.data.Options()

# Stop magic stuff that eats up RAM:
options.autotune.enabled = False
options.experimental_distribute.auto_shard_policy = (
  tf.data.experimental.AutoShardPolicy.OFF)
options.experimental_optimization.inject_prefetch = False

data = data.with_options(options)
```

---

### Splits

# Splits and slicing

All TFDS datasets expose various data splits (e.g. `'train'`, `'test'`) which
can be explored in the
[catalog](https://www.tensorflow.org/datasets/catalog/overview). Any
alphabetical string can be used as split name, apart from `all` (which is a
reserved term which corresponds to the union of all splits, see below).

In addition of the "official" dataset splits, TFDS allow to select slice(s) of
split(s) and various combinations.

## Slicing API

Slicing instructions are specified in `tfds.load` or
`tfds.DatasetBuilder.as_dataset` through the `split=` kwarg.

```python
ds = tfds.load('my_dataset', split='train[:75%]')
```

```python
builder = tfds.builder('my_dataset')
ds = builder.as_dataset(split='test+train[:75%]')
```

Split can be:

*   **Plain split names** (a string such as `'train'`, `'test'`, ...): All
    examples within the split selected.
*   **Slices**: Slices have the same semantic as
    [python slice notation](https://docs.python.org/3/library/stdtypes.html#common-sequence-operations).
    Slices can be:
    *   **Absolute** (`'train[123:450]'`, `train[:4000]`): (see note below for
        caveat about read order)
    *   **Percent** (`'train[:75%]'`, `'train[25%:75%]'`): Divide the full data
        into even slices. If the data is not evenly divisible, some percent
        might contain additional examples. Fractional percent are supported.
    *   **Shard** (`train[:4shard]`, `train[4shard]`): Select all examples in
        the requested shard. (see `info.splits['train'].num_shards` to get the
        number of shards of the split)
*   **Union of splits** (`'train+test'`, `'train[:25%]+test'`): Splits will be
    interleaved together.
*   **Full dataset** (`'all'`): `'all'` is a special split name corresponding to
    the union of all splits (equivalent to `'train+test+...'`).
*   **List of splits** (`['train', 'test']`): Multiple `tf.data.Dataset` are
    returned separately:

```python
# Returns both train and test split separately
train_ds, test_ds = tfds.load('mnist', split=['train', 'test[:50%]'])
```

Note: Due to the shards being
[interleaved](https://www.tensorflow.org/api_docs/python/tf/data/Dataset?version=nightly#interleave),
order isn't guaranteed to be consistent between sub-splits. In other words
reading `test[0:100]` followed by `test[100:200]` may yield examples in a
different order than reading `test[:200]`. See
[determinism guide](https://www.tensorflow.org/datasets/determinism#determinism_when_reading)
to understand the order in which TFDS read examples.

## `tfds.even_splits` & multi-host training

`tfds.even_splits` generates a list of non-overlapping sub-splits of the same
size.

```python
# Divide the dataset into 3 even parts, each containing 1/3 of the data
split0, split1, split2 = tfds.even_splits('train', n=3)

ds = tfds.load('my_dataset', split=split2)
```

This can be particularly useful when training in a distributed setting, where
each host should receive a slice of the original data.

With `Jax`, this can be simplified even further using
`tfds.split_for_jax_process`:

```python
split = tfds.split_for_jax_process('train', drop_remainder=True)
ds = tfds.load('my_dataset', split=split)
```

`tfds.split_for_jax_process` is a simple alias for:

```python
# The current `process_index` loads only `1 / process_count` of the data.
splits = tfds.even_splits('train', n=jax.process_count(), drop_remainder=True)
split = splits[jax.process_index()]
```

`tfds.even_splits`, `tfds.split_for_jax_process` accepts on any split value as
input (e.g. `'train[75%:]+test'`)

## Slicing and metadata

It is possible to get additional info on the splits/subsplits (`num_examples`,
`file_instructions`,...) using the
[dataset info](https://www.tensorflow.org/datasets/overview#access_the_dataset_metadata):

```python
builder = tfds.builder('my_dataset')
builder.info.splits['train'].num_examples  # 10_000
builder.info.splits['train[:75%]'].num_examples  # 7_500 (also works with slices)
builder.info.splits.keys()  # ['train', 'test']
```

## Cross validation

Examples of 10-fold cross-validation using the string API:

```python
vals_ds = tfds.load('mnist', split=[
    f'train[{k}%:{k+10}%]' for k in range(0, 100, 10)
])
trains_ds = tfds.load('mnist', split=[
    f'train[:{k}%]+train[{k+10}%:]' for k in range(0, 100, 10)
])
```

The validation datasets are each going to be 10%: `[0%:10%]`, `[10%:20%]`, ...,
`[90%:100%]`. And the training datasets are each going to be the complementary
90%: `[10%:100%]` (for a corresponding validation set of `[0%:10%]`), `[0%:10%]
+ [20%:100%]`(for a validation set of `[10%:20%]`),...

## `tfds.core.ReadInstruction` and rounding

Rather than `str`, it is possible to pass splits as `tfds.core.ReadInstruction`:

For example, `split = 'train[50%:75%] + test'` is equivalent to:

```python
split = (
    tfds.core.ReadInstruction(
        'train',
        from_=50,
        to=75,
        unit='%',
    )
    + tfds.core.ReadInstruction('test')
)
ds = tfds.load('my_dataset', split=split)
```

`unit` can be:

*   `abs`: Absolute slicing
*   `%`: Percent slicing
*   `shard`: Shard slicing

`tfds.ReadInstruction` also has a rounding argument. If the number of example in
the dataset is not divide evenly:

*   `rounding='closest'` (default): The remaining examples are distributed among
    the percent, so some percent might contain additional examples.
*   `rounding='pct1_dropremainder'`: The remaining examples are dropped, but
    this guarantee all percent contain the exact same number of example (eg:
    `len(5%) == 5 * len(1%)`).

### Reproducibility & determinism

During generation, for a given dataset version, TFDS guarantee that examples are
deterministically shuffled on disk. So generating the dataset twice (in 2
different computers) won't change the example order.

Similarly, the subsplit API will always select the same `set` of examples,
regardless of platform, architecture, etc. This mean `set('train[:20%]') ==
set('train[:10%]') + set('train[10%:20%]')`.

However, the order in which examples are read might **not** be deterministic.
This depends on other parameters (e.g. whether `shuffle_files=True`).

---