TensorFlow Datasets

TensorFlow Datasets (TFDS) is a collection of hundreds of public datasets packaged behind a single tfds.load() call, handling the download, checksum verification, and on-disk conversion into a standard, versioned format so every user gets identical, reproducible splits. Rather than each project writing its own scraping and parsing code for MNIST, ImageNet, or a text corpus, TFDS centralizes that logic once per dataset and exposes the result as a tf.data.Dataset, a NumPy array via as_numpy(), or a random-access data source for JAX training loops.

Under the hood, each dataset is defined by a DatasetBuilder subclass that declares its features, splits, and download sources; TFDS then handles sharding the prepared data into TFRecord, ArrayRecord, Parquet, or Riegeli files depending on the desired access pattern (sequential streaming vs. random access). This separation between a small builder definition and a shared execution engine is what lets the project scale to hundreds of datasets across image, text, audio, video, and RL domains while keeping the public API — tfds.load(name, split=...) — identical across all of them.

What You Get

• A single tfds.load() / tfds.builder() API that works identically across hundreds of pre-defined datasets
• Automatic download, checksum verification, and one-time preparation into an efficient sharded format
• Output as a tf.data.Dataset, a NumPy array via as_numpy(), or a random-access data source for JAX/PyTorch-style loops
• Deterministic, versioned dataset splits so results are reproducible across runs and machines
• Support for multiple on-disk storage formats (TFRecord, ArrayRecord, Parquet, Riegeli) chosen per access pattern
• A DatasetBuilder framework for registering and contributing new datasets with the same guarantees

Common Use Cases

• Loading standard benchmark datasets (MNIST, ImageNet, CIFAR, GLUE, etc.) for model training and evaluation without writing custom parsers
• Building JAX training pipelines that need a random-access data source rather than a streaming iterator
• Reproducing published results where exact, versioned dataset splits matter
• Prototyping on a new dataset by wrapping it in a DatasetBuilder and getting sharding, checksums, and splits for free
• Feeding tf.data input pipelines for large-scale TensorFlow training jobs

Under The Hood

Architecture TFDS is organized around a DatasetBuilder abstraction (tensorflow_datasets/core/dataset_builder.py): each concrete dataset (under top-level domain folders like image/, text/, audio/, rl_unplugged/) subclasses GeneratorBasedBuilder or BeamBasedBuilder and declares its features, split generators, and download URLs, while the shared DatasetBuilder base class in core/ handles the generic download-checksum-shard-serialize pipeline (core/download/, core/example_serializer.py, core/shuffle.py, core/file_adapters.py). core/registered.py and core/load.py form a registry/lookup layer that resolves a dataset name string to its builder class, and core/read_only_builder.py lets a prepared dataset be loaded purely from its serialized directory without re-importing the original builder code. This is a layered, plugin-style architecture: adding a dataset means writing a builder subclass, not touching the shared engine, and the core pipeline is what breaks if it changes — dataset-specific code is comparatively isolated.

Tech Stack The library targets Python 3.10+ and is built with setuptools/pyproject.toml. Core runtime dependencies include absl-py, etils (for path/config utilities), numpy, protobuf, tensorflow-metadata, dm-tree, pyarrow, and array_record, reflecting its role as glue between raw data formats and multiple ML frameworks rather than a single-framework tool. Storage backends are pluggable via file_adapters.py, which maps a FileFormat enum (TFRecord, ArrayRecord, Parquet, Riegeli) to matching reader/writer adapter classes. CI is run through GitHub Actions workflows (pytest.yml, pytest-template.yml) driving pytest, and packaging optionally builds a tfds-nightly variant driven by an environment-variable-based version suffix in setup.py.

Code Quality Tests are colocated with source as *_test.py files throughout core/ (e.g. dataset_builder_test.py, file_adapters_test.py, registered_test.py, shuffle_test.py), giving the core pipeline dense, file-level test coverage exercised via pytest in CI. The codebase is extensively type-annotated (from __future__ import annotations, Optional/Type/Sequence hints throughout dataset_builder.py and load.py) and ships a .pylintrc, with pylint listed as a dev dependency. Comment density in core modules is high, with docstrings and inline rationale accompanying non-obvious logic such as format-specific sharding rules.

What Makes It Unique TFDS’s distinguishing design choice is treating dataset preparation as a first-class, versioned build step rather than an ad hoc download script: every dataset gets deterministic splits, checksum-verified downloads, and a choice of serialization format (including ArrayRecord and Parquet for random-access reads, which most streaming-only dataset libraries don’t offer), letting the same catalog serve both tf.data streaming pipelines and JAX-style random-access training loops from one builder definition.