Crate webdataset

Expand description

High-performance sequential dataset loading, in the WebDataset format.

A WebDataset is a set of tar archives (“shards”). Inside a shard, the files that share a basename make up one training sample:

imagenet-000000.tar
  n03991062_24866.jpg   n03991062_24866.cls
  n03995372_9042.jpg    n03995372_9042.cls

Reading is purely sequential, which is what makes the format fast: no seeks, no index, no per-file requests. A shard reads at the full bandwidth of the device or the network link, and a dataset is just a list of URLs — nothing needs to be registered, converted, or mounted first.

§Reading a dataset

use webdataset::{Decoder, WebDataset};
use webdataset::filters::{SampleIteratorExt, TupleIteratorExt};

let dataset = WebDataset::builder("https://host/imagenet-{000000..000146}.tar")
    .shard_shuffle(100)
    .build()?
    .shuffle(1000)
    .decode(Decoder::default());

for batch in dataset.iter().to_tuple(["jpg;png", "cls"]).batched(64, true) {
    let batch = batch?;
    let images = &batch[0];
    let labels = &batch[1];
}

§Writing a dataset

use std::sync::Arc;
use webdataset::encode::DefaultEncoder;
use webdataset::{Sample, ShardWriter, Value};

let mut writer = ShardWriter::new(pattern)?
    .with_encoder(Arc::new(DefaultEncoder::new()))
    .with_max_count(10_000);

for i in 0..100 {
    let mut sample = Sample::with_key(format!("sample{i:06}"));
    sample.insert("cls", Value::Int(i % 10));
    sample.insert("txt", Value::Text(format!("sample number {i}")));
    writer.write(&sample)?;
}
writer.close()?;

§How a pipeline is put together

WebDataset assembles the usual stages for you, but the pieces are public and a pipeline can be built by hand:

SimpleShardList      a list of shard URLs
  -> SplitByNode     keep this rank's shards
  -> SplitByWorker   keep this worker's shards
  -> Shuffle         shuffle the shard order
  -> ShardsToSamples open each shard, group its files into samples
  -> Shuffle         shuffle samples within a buffer
  -> Decode          turn bytes into images, tensors, JSON

Everything up to this point is a Stage producing Samples, so it can be re-run each epoch. Transformations that change the item type — projecting to tuples, batching — are ordinary iterator adapters from filters, applied to WebDataset::iter().

§Scaling out

DataLoader runs one copy of the pipeline per worker thread. Shards are divided between workers by the SplitByWorker stage, and between distributed processes by SplitByNode. When exact partitioning is awkward — many nodes, few shards — use resampled(true) instead and let each worker draw shards with replacement.

§Errors

Every stream yields Result<Sample>. What happens when a shard is corrupt or a field fails to decode is decided by a Handler: forward the error, drop the sample, or end the stream. Streaming a petabyte means meeting some bad bytes, so warn_and_continue is a common choice.

§Features

feature	adds
`threads` (default)	per-shard read-ahead and the multi-worker loader
`subprocess` (default)	`pipe:` and the `curl`/`gsutil`/`ais` schemes
`yaml` (default)	multi-source dataset specifications
`image`	`.jpg`, `.png` and friends, via the `image` crate
`msgpack`	`.mp` and `.msg`
`cbor`	`.cbor`
`npz`	NumPy `.npz` archives
`zstd`, `bzip2`, `xz`	shards in those containers
`async`	read shards from any `AsyncRead`, yielding a `Stream`
`wasm-js`	host randomness on `wasm32-unknown-unknown`
`full`	every format, plus threads, subprocesses and async

.npy, .ten, .json, .txt, .cls and gzip need no features.

§Async

With the async feature, asynch mirrors everything above: the same archive parser, the same decoders, the same shuffling and batching, driven by futures and producing a Stream rather than an Iterator. Reach for it when shards arrive over a network — a blocking reader holds a thread for the whole of a transfer, an async one holds only a task, and raising concurrency overlaps the fetches.

Only the I/O is async; decoding and batching are the same code the blocking pipeline runs, which is why the two read identically — a property the test suite checks shard by shard. See the asynch module for a worked example.

§WebAssembly

WebAssembly has no threads to spawn and no processes to run, so turn both features off and hand the shard bytes over yourself with MemoryOpener. Everything above the transport — the archive parser, the decoders, shuffling, batching — is unchanged.

use std::sync::Arc;
use webdataset::pipeline::DataPipeline;
use webdataset::shardlists::SimpleShardList;
use webdataset::sources::{MemoryOpener, ShardsToSamples};
use webdataset::stages::Decode;

let opener = MemoryOpener::new().with("mem://shard-000.tar", bytes);
let dataset = DataPipeline::new()
    .with(SimpleShardList::verbatim(["mem://shard-000.tar"]))
    .with(ShardsToSamples::new(Arc::new(opener)))
    .with(Decode::basic());

assert_eq!(dataset.iter().count(), 90);

For any other transport, register a scheme with webdataset_io::register_scheme.

webdataset-core and webdataset-tenbin go further and build without the standard library at all; see their documentation.

Re-exports§

pub use batch::collate_samples;
pub use batch::collate_tuples;
pub use batch::uncollate_sample;
pub use dataset::NodeSplit;
pub use dataset::WebDataset;
pub use dataset::WebDatasetBuilder;
pub use decode::DecodeHandler;
pub use decode::Decoded;
pub use decode::Decoder;
pub use encode::DefaultEncoder;
pub use filters::SampleIteratorExt;
pub use filters::TupleIteratorExt;
pub use loader::DataLoader;
pub use mix::RandomMix;
pub use mix::RoundRobin;
pub use pipeline::DataPipeline;
pub use pipeline::SampleStream;
pub use pipeline::Stage;
pub use shardlists::ResampledShards;
pub use shardlists::SimpleShardList;
pub use shardlists::SplitByNode;
pub use shardlists::SplitByWorker;
pub use sources::MemoryOpener;
pub use sources::Opener;
pub use sources::ShardsToSamples;
pub use images::ImageHandler;image
pub use images::ImageSpec;image
pub use webdataset_core as core;
pub use webdataset_io as io;
pub use webdataset_shard as shard;
pub use webdataset_tenbin as tenbin;

Modules§

asynchasync: Asynchronous pipelines.
batch: Collating samples into batches.
braceexpand: Bash-style brace expansion.
dataset: The high-level entry point: WebDataset.
decode: Turning raw file bytes into usable values.
encode: Turning values back into the bytes stored in a shard.
filters: Iterator adapters for sample streams.
gopen: Open a URL, whatever its scheme.
handlers: Pluggable exception handlers.
imagesimage: Image decoding and encoding.
loader: Running a pipeline over several worker threads.
mix: Combining several datasets into one stream.
pipeline: Composable sample pipelines.
prelude: The imports a typical pipeline needs, in one line.
shardlists: Producing the stream of shard URLs a pipeline reads from.
sources: Turning shard URLs into samples.
stages: Stage wrappers around the filters.

Structs§

FileCache: Downloads shards to a local directory and serves them from there.
Sample: One training example: an insertion-ordered map from file extension to value.
Selection: Chooses which archive members to read, and under what name.
ShardWriter: Writes samples across a numbered series of shards.
TarWriter: Writes samples into a single tar archive.
Tensor: A dense, C-ordered array of numbers.
WorkerInfo: Which process and worker this code is running in.

Enums§

Action: What a pipeline stage should do after an error.
DType: The element type of a Tensor.
Error: The error type used throughout the webdataset crates.
Value: A value carried by a sample field.

Traits§

Handler: Decides how a pipeline stage reacts to an error.

Functions§

braceexpand: Expand a brace expression into the list of strings it denotes.
gopen: Open url for reading.
gopen_write: Open url for writing.

Type Aliases§

HandlerRef: A shared, cheaply clonable handler.
Result: Convenience alias for results produced by this workspace.