Repository: pytorch
Issue: [RFC] Add tar-based IterableDataset implementation to PyTorch
Labels: feature,module: dataloader,triaged

Problem
=======

As datasets become larger and larger, storing training samples as individual files becomes impractical and inefficient. This can be addressed using sequential storage formats and sharding (see "Alternatives Considered" for other implementations).  PyTorch lacks such a common storage format right now.


Proposal
========

[WebDataset](http://github.com/tmbdev/webdataset) provides an implementation of `IterableDataset` based on sharded tar archives. This format provides efficient I/O for very large datasets, makes migration from file-based I/O easy, and works well locally, with cloud storage, and web servers.  It also provides a simple, standard format in which large datasets can be distributed easily and used directly without unpacking.

The implementation is small (1000-1800 LOC) and has no external dependencies.  The proposal is to incorporate the `webdataset.Dataset` class into the PyTorch base distribution. The source repository is here:

http://github.com/tmbdev/webdataset

My suggestion would be to incorporate the library with minimal changes into its own subpackage. I can perform the integration and generate a PR once there is general agreement.


More Background
===============

For a general introduction to how we handle large scale training with WebDataset, see [this YouTube playlist](https://www.youtube.com/playlist?list=PL0dsKxFNMcX4XcB0w1Wm-pvSfQu-eWM26)

The WebDataset library (github.com/tmbdev/webdataset) provides an implementation of `IterableDataset` that uses POSIX tar archives as its native storage format.  The format itself is based on a simple convention:

- datasets are POSIX tar archives
- each training sample consists of adjacent files with the same basename
- shards are numbered consecutively

For example, ImageNet is stored in 147 separate 1 Gbyte shards with names `imagenet-train-0000.tar` to `imagenet-train-0147.tar`; the contents of the first shard are:

```
-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n03991062_24866.cls
-r--r--r-- bigdata/bigdata 108611 2020-05-08 21:23 n03991062_24866.jpg
-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n07749582_9506.cls
-r--r--r-- bigdata/bigdata 129044 2020-05-08 21:23 n07749582_9506.jpg
-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n03425413_23604.cls
-r--r--r-- bigdata/bigdata 106255 2020-05-08 21:23 n03425413_23604.jpg
-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n02795169_27274.cls
```

Datasets in WebDataset format can be used directly after downloading without unpacking; they can also be mounted as a file system.  Content in WebDataset format can used any file-based compression scheme.  In addition, the tar file itself can also be compressed and WebDataset will transparently decompress it.

WebDatsets can be used directly from local disk, from web servers (hence the name), from cloud storage, and from object stores, just by changing a URL.

WebDataset readers/writers are easy to implement (we have Python, Golang, and C++ implementations).

WebDataset performs shuffling both at the shard level and at the sample level.  Splitting of data across multiple workers is performed at the shard level using a user-provided `shard_selection` function that defaults to a function that splits based on `get_worker_info`. (WebDataset can be combined with the `tensorcom` library to offload decompression/data augmentation and provide RDMA and direct-to-GPU loading; see below.)

We are storing and processing petabytes of training data as tar archives and are using the format for unsupervised learning from video, OCR and scene text recognition, and large scale object recognition experiments.  The same code and I/O pipelines work efficiently on the desktop, on a local cluster, or in the cloud.  [Our benchmarks](https://arxiv.org/abs/2001.01858) show scalability and the ability to take advantage of the full I/O bandwidth of each disk drive across very large training jobs and storage clusters.


Code Sample
===========

This shows how to use WebDataset with ImageNet.

```Python
import webdataset as wds
improt ...

shardurl = "/imagenet/imagenet-train-{0000..0147}.tar"

normalize = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225])

preproc = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    normalize,
])

dataset = (
    wds.Dataset(shardurl)
    .shuffle(1000)
    .decode("pil")
    .rename(image="jpg;png", data="json")
    .map_dict(image=preproc)
    .to_tuple("image", "data")
)

loader = torch.utils.data.DataLoader(dataset, batch_size=64, num_workers=8)

for inputs, targets in loader:
    ...
```

WebDataset uses a fluent API for configuration that internally builds up a processing pipeline. Without any added processing stages, WebDataset just iterates through each training sample as a dictionary:

```Python
# load from a web server using a separate client process
shardurl = "pipe:curl -s http://server/imagenet/imagenet-train-{0000..0147}.tar"

dataset = wds.Dataset(shardurl)

for sample in dataset:
    # sample["jpg"] contains the raw image data
    # sample["cls"] contains the class
    ...
```


Alternatives Considered
=======================

Keeping WebDataset as as Separate Library
-----------------------------------------

WebDataset is perfectly usable as a separate library, so why not keep it that way?

- Users of deep learning libraries expect an efficient data format that avoids the "many small file" problem; Tensorflow provides TFRecord/tf.Example Making WebDataset part of PyTorch itself provides a straightforward solution for most users and reduces dependencies.
- Providing a format reader in the standard library encourages dataset standardization, tool building, and adoption of the same format by other libraries.
- Having a realistic implementation of `IterableDataset` in PyTorch provides a common reference against which to address issues and provide test cases in the `DataLoader` implementation and its use of `IterableDataset`

Many of the larger datasets distributed in `torchvision` could be distributed easily in WebDataset format. This would allow users to either train directly against web-hosted data, or to train on the datasets immediately after downloading without unpacking. The way data is arranged in WebDataset also allows users to download just a few shards for testing code locally, and then use the entire dataset when running on a cluster. Furthermore, the way `webdataset.Dataset` works in most cases, no special code is needed in order to read them; many training jobs can be retargeted to different datasets simply by using a different URL for the dataset.


TFRecord+protobuf, Parquet
----------------------------

These formats are suitable for large scale data processing for machine learning and deep learning applications and some datasets exist in this format and more will continue to be generated for the Tensorflow ecosystem. However, they are not good candidates for incorporating into PyTorch as core feature because:

- TFRecord+protobuf (tf.Example) and Parquet have significant external library dependencies; in contrast, WebDataset is pure Python, using the built-in libraries for tar decoding.
- Protobuf and Parquet represent serialized data structures, while WebDataset uses file-based representations. File based representations make creation of WebDatasets easy, result in bit-identical representations of data in WebDatasets, and allow existing file-based codecs and I/O pipelines to be reused easily.
- Except for their native frameworks, there is not much third party support for these formats; in contrast, POSIX tar archive libraries exist in all major programming languages and are easily processed, either by unpacking shards or using tools like `tarp`.

Note that WebDataset supports usage scenarios similar to TFRecord+protobuf, since serialized data structures can be incorporated as files, and WebDataset will decode them automatically. For example, OpenImages multi-instance data is simply stored in a `.json` file accompanying each `.jpg` file:

```
$ curl -s http://storage.googleapis.com/nvdata-openimages/openimages-train-000554.tar | tar tf -
0b946d1d5201c06c.jpg
0b946d1d5201c06c.json
05c9d72ff9e64010.jpg
05c9d72ff9e64010.json
...
$
```

zip instead of tar
--------------------

The zip format is another archival format. Unlike tar format, which is just a sequence of records, zip format stores a file index at the very end of the file, making it unsuitable for streaming. Tar files can be made random access (and, in fact, can be mounted as file systems), but they use a separate index file to support that functionality.

LMDB, HDF5, Databases
---------------------

These formats are not suitable for streaming and require the entire dataset to fit onto local disk. In addition, while they nominally solve the "many small file problems", they don't solve the problem that indexing into the dataset still results in expensive seek operations.

Local File System Caches
------------------------

An approach for extending file-system based I/O pipelines to large distributed storage systems is to use some form of "pre-caching" or "staging" on a local NVMe drive. Generally, there is little advantage to this. For large datasets, it does not increase throughput. Input pipelines still need to be modified to schedule the pre-caching. And generally, this requires volume plugins or virtual file system support. A similar effect can be achieved with WebDataset by simply unpacking shards to the local file system when direct file access is required.


Related Software
================

[AIStore](http://github.com/NVIDIA/AIStore) is an open source object store capable of full bandwidth disk-to-GPU data delivery (meaning that if you have 1000 rotational drives with 200 MB/s read speed, AIStore actually delivers an aggregate bandwidth of 200 GB/s to the GPUs). AIStore is fully compatible with WebDataset as a client, and in addition understands the WebDataset format, permitting it to perform shuffling, sorting, ETL, and some map-reduce operations directly in the storage system.

[tarp](http://github.com/tmbdev/tarp) is a small command line program for splitting, merging, shuffling, and processing tar archives and WebDataset datasets.

[tensorcom](http://github.com/NVLabs/tensorcom) is a library supporting distributed data augmentation and RDMA to GPU.

[webdataset-examples](http://github.com/tmbdev/webdataset-examples) contains an example (and soon more examples) of how to use WebDataset in practice.

[Bigdata 2019 Paper with Benchmarks](https://arxiv.org/abs/2001.01858)

cc @SsnL