Add Dataset Copy Section

Writing custom datasets

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

Check our list of datasets to see if the dataset you want
is already present.

TL;DR Copy Chapter

The easiest way to write a new dataset is to use the
TFDS CLI:

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:

import my.project.datasets.my_dataset  # Register my_dataset
ds = tfds.load('my_dataset')  # my_dataset registered

Overview Copy Chapter

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 Copy Chapter

Default template: tfds new tfds new')" title="Copy section prompt for LLMs"> Copy Section

Use TFDS CLI to generate the required
template python files.

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:

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 Copy Section

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, see our
huge dataset guide)

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

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
to take care of most data processing.

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

_info: dataset metadata _info: dataset metadata')" title="Copy section prompt for LLMs"> Copy Section

_info returns the tfds.core.DatasetInfo containing the
dataset metadata.

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
or the
feature connector guide for
more info.
* disable_shuffling: See section
Maintain dataset order.

Writing the BibText CITATIONS.bib file:

* Search the dataset website for citation instruction (use that in BibTex
format).
* For arXiv papers: find the paper and click the
BibText link on the right-hand side.
* Find the paper on Google Scholar 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 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.
* 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.

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 _split_generators: downloads and splits data')" title="Copy section prompt for LLMs"> Copy Section

#### 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), which are
pathlib.Path-like objects.

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

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:

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.

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 (e.g.
split='train[80%:]'). Note that any alphabetical string can be used as split
name, apart from the aforementioned all.

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: Example generator')" title="Copy section prompt for LLMs"> Copy Section

_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). 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
for more info.

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 objects directly
compatible with Google Cloud storage:

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.

#### 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.
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
and to the
LazyImportsTest.
* 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) Copy Section

Some datasets may have multiple variants, or options for how the data is
preprocessed and written to disk. For example,
cycle_gan has one
config per object pairs (cycle_gan/horse2zebra, cycle_gan/monet2photo,...).

This is done through tfds.core.BuilderConfigs:

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

@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
MyDatasetConfigs that the dataset exposes.

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 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
for an example of a dataset that uses BuilderConfigs.

Version Copy Section

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.

Add an import for registration Copy Section

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

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.

Check for common implementation gotchas Copy Section

Please check for the
common implementation gotchas.

Test your dataset Copy Chapter

Download and prepare: tfds build tfds build')" title="Copy section prompt for LLMs"> Copy Section

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

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
for full list of flags.

Checksums Copy Section

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 Copy Section

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).
* 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.

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.

python my_dataset_test.py

Send us feedback Copy Chapter

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.

---

Add Dataset Collection Copy Section

Add a new dataset collection

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

Overview Copy Chapter

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

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 Copy Chapter

All dataset collections are implemented subclasses of
tfds.core.dataset_collection_builder.DatasetCollection.

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

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 info: dataset collection metadata')" title="Copy section prompt for LLMs"> Copy Section

The info method returns the
dataset_collection_builder.DatasetCollectionInfo
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 datasets: define the datasets in the collection')" title="Copy section prompt for LLMs"> Copy Section

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.
For example:

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
method provides a more compact way to express the same as above:

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 Copy Chapter

DatasetCollectionTestBase
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:

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.

python my_dataset_test.py

Feedback Copy Chapter

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.

---

Announce Proxy Copy Section

Introducing TensorFlow Datasets

TensorFlow Datasets is now released on PyPI:

pip install tensorflow-datasets

Read the blog post
to learn more.

---

Beam Datasets Copy Section

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.

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 Copy Chapter

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,
make sure to specify the dataset config you want to generate or it will default
to generate all existing configs. For example, for
wikipedia, use tfds
build wikipedia/20200301.en instead of tfds build wikipedia.

On Google Cloud Dataflow Copy Section

To run the pipeline using
Google Cloud Dataflow and take advantage
of distributed computation, first follow the
Quickstart instructions.

Once your environment is set up, you can run the
tfds build CLI
using a data directory on GCS and
specifying the
required options
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:

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:

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.

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.

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

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 Copy Section

To run your script locally using the
default Apache Beam runner
(it must fit all data in memory), the command is the same as for other datasets:

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 for a list of supported
runtimes.

With Apache Flink Copy Section

To run the pipeline using Apache Flink you can read
the
official documentation.
Make sure your Beam is compliant with
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:

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:

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 Copy Section

To generate the dataset on Beam, the API is the same as for other datasets. You
can customize the
beam.Pipeline
using the beam_options (and beam_runner) arguments of DownloadConfig.

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 Copy Chapter

Prerequisites Copy Section

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

* Be familiar with the
tfds dataset creation guide
as most of the content still applies for Beam datasets.
* Get an introduction to Apache Beam with the
Beam programming guide.
* If you want to generate your dataset using Cloud Dataflow, read the
Google Cloud Documentation
and the
Apache Beam dependency guide.

Instructions Copy Section

If you are familiar with the
dataset creation guide,
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:

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

Beam dataset:

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 Copy Section

Here is an example of a Beam dataset.

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 Copy Section

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

tfds build my_dataset --register_checksums

Pipeline using TFDS as input Copy Chapter

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

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 Copy Section

Common implementation gotchas

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

Legacy SplitGenerator should be avoided SplitGenerator should be avoided')" title="Copy chapter prompt for LLMs"> Copy Chapter

The old tfds.core.SplitGenerator API is deprecated.

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:

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 Copy Chapter

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
(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 Copy Chapter

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

_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 Copy Chapter

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

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
tfds.show_examples, tfds.as_dataframe

Forgot image shape Copy Chapter

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

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 tfds.features.Tensor')" title="Copy chapter prompt for LLMs"> Copy Chapter

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 Copy Chapter

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

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 Copy Chapter

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

_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 should be used
instead to let users dynamically create the subsplits they want:

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

Python style guide Copy Chapter

Prefer to use pathlib API Copy Section

Instead of the tf.io.gfile API, it is preferable to use the
pathlib API. All dl_manager
methods returns pathlib-like objects compatible with GCS, S3,...

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 self, it should be a function')" title="Copy section prompt for LLMs"> Copy Section

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:

x = f(y)  # Clear inputs/outputs
x = self.f(y)  # Does f depend on additional hidden variables ? Is it stateful ?

Lazy imports in Python Copy Chapter

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.

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
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:

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

---

Contribute Copy Section

Contribute to the TFDS repository

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

Get started Copy Chapter

* If you're new with TFDS, the easiest way to get started is to implement one
of our
requested datasets,
focusing on the most requested ones.
Follow our guide 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.
Starts with the ones labeled with
contribution-welcome
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, and
comment on issues you're working on.

Setup Copy Chapter

Cloning the repo Copy Section

To get started, clone or download the
Tensorflow Datasets repository and
install the repo locally.

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

Install the development dependencies:

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 Copy Section

When developing with Visual Studio Code, our
repo comes with some
pre-defined settings
to help development (correct indentation, pylint,...).

Note: enabling test discovery in VS Code may fails due to some VS Code bugs
#13301 and
#6594. To solve the
issues, you can look at the test discovery logs:

* If you are encountering some TensorFlow warning message, try
this fix.
* If discovery fail due to missing import which should have been installed,
please send a PR to update the dev pip install.

PR checklist Copy Chapter

Sign the CLA Copy Section

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 Copy Section

* 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

Check your style guide Copy Section

Our style is based on
Google Python Style Guide,
which is based on
PEP 8 Python style guide. New code
should try to follow
Black code style
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:

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.

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 Copy Section

Classes and functions should be documented with docstrings and typing
annotation. Docstrings should follow the
Google style.
For example:

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 Copy Section

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

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
if you're not familiar with pytest.

Tests for DatasetBuilders are special and are documented in the
guide to add a dataset.

Send the PR for reviews! Copy Section

Congrats! See
GitHub Help for more
information on using pull requests.

---

Datasets Versioning Copy Section

Datasets versioning

Definition Copy Chapter

Versioning can refer to different meaning:

* The TFDS API version (pip version): tfds.__version__
* The public dataset version, independent from TFDS (e.g.
Voc2007,
Voc2012). In TFDS each public dataset version should be implemented as an
independent dataset:
* Either through
builder configs:
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 Copy Chapter

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+).

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 Copy Chapter

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

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:
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 Copy Chapter

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

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 Copy Chapter

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:

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 Copy Chapter

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:

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:

class MNIST(tfds.core.GeneratorBasedBuilder):
  ...
  def _generate_examples(self, path):
    if self.info.version >= '2.0.0':
      ...
    else:
      ...

---

Decode Copy Section

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 Copy Chapter

Skipping the image decoding Copy Section

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.

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 Copy Section

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.

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 Copy Section

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.

@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:

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 Copy Section

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.

@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:

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. Copy Section

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

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:

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 Copy Section

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 Copy Chapter

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
look as follows:

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

Add metadata Copy Chapter

Provide the feature structure Copy Section

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:

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:

{
    '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:

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)

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

Get statistics on splits Copy Section

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

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

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 Copy Section

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

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 Copy Chapter

Directly from folder Copy Section

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):

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 Copy Section

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):

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) Copy Section

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/:

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

---

Features Copy Section

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 Copy Chapter

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

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
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):

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:

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 Copy Chapter

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:

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:

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

Access metadata Copy Chapter

See the
introduction doc
to access features metadata (label names, shape, dtype,...). Example:

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 tfds.features.FeatureConnector')" title="Copy chapter prompt for LLMs"> Copy Chapter

If you believe a feature is missing from the
available features,
please open a new issue.

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
for an example.

Note: Make sure to test your Feature connectors with self.assertFeature and
tfds.testing.FeatureExpectationItem. Have a look at
test examples:

For more info, have a look at tfds.features.FeatureConnector documentation.
It's also best to look at
real examples.

---

Format Specific Dataset Builders Copy Section

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
which take care of most data processing for a specific data format.

Datasets based on tf.data.Dataset tf.data.Dataset')" title="Copy chapter prompt for LLMs"> Copy Chapter

If you want to create a TFDS dataset from a dataset that's in tf.data.Dataset
(reference)
format, then you can use tfds.dataset_builders.TfDataBuilder (see
API docs).

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 Copy Section

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.

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.Datasets
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:

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.

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 Copy Section

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

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 Copy Chapter

The format Copy Section

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 CroissantBuilder')" title="Copy section prompt for LLMs"> Copy Section

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:

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 Copy Chapter

The format Copy Section

CoNLL 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
dataset, which follows the CoNLL annotation format:

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 ConllDatasetBuilder')" title="Copy section prompt for LLMs"> Copy Section

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:

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
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 Copy Section

Consider
conll2003
as an example of a dataset implemented using the CoNLL-specific dataset builder.

CLI Copy Section

The easiest way to write a new CoNLL-based dataset is to use the
TFDS CLI:

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

CoNLL-U Copy Chapter

The format Copy Section

CoNLL-U 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.
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:

* 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:

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

ConllUDatasetBuilder ConllUDatasetBuilder')" title="Copy section prompt for LLMs"> Copy Section

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:

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
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 -
you might need to update the process_example_fn attribute of your
generate_examples
function (see the
xtreme_pos
daset as an example).

Examples Copy Section

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

* universal_dependencies
* xtreme_pos

CLI Copy Section

The easiest way to write a new CoNLL-U based dataset is to use the
TFDS CLI:

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

---

Gcs Copy Section

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 Copy Chapter

Some datasets are available directly in our GCS bucket
gs://tfds-data/datasets/
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 Copy Chapter

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

Simplified instructions Copy Section

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

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:

gcloud auth application-default login

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

export GOOGLE_APPLICATION_CREDENTIALS=<JSON_FILE_PATH>

Using Google Cloud Storage to store preprocessed data Copy Chapter

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) 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? Copy Section

Create a GCS bucket
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"

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

Caveats: Copy Section

* 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 Copy Chapter

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 Copy Section

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
still applies.

Benchmark datasets Copy Chapter

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

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) Copy Chapter

All TFDS datasets store the data on disk in the
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 Copy Section

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

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 Copy Section

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 Copy Section

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.

(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 Copy Chapter

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

Shuffle and training Copy Section

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.

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

Additionally, when shuffle_files=True, TFDS disables
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) Copy Section

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.

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

Auto-shard your data across workers (Jax) Copy Section

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.

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:

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 Copy Section

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.

Skip unused features Copy Section

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! Copy Chapter

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 Copy Section

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 Copy Section

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

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 Copy Section

Splits and slicing

All TFDS datasets expose various data splits (e.g. 'train', 'test') which
can be explored in the
catalog. 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 Copy Chapter

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

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

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

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,
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
to understand the order in which TFDS read examples.

tfds.even_splits & multi-host training tfds.even_splits & multi-host training')" title="Copy chapter prompt for LLMs"> Copy Chapter

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

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:

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:

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 Copy Chapter

It is possible to get additional info on the splits/subsplits (num_examples,
file_instructions,...) using the
dataset info:

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 Copy Chapter

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

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 tfds.core.ReadInstruction and rounding')" title="Copy chapter prompt for LLMs"> Copy Chapter

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

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

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 Copy Section

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

---