WebDataset

Warnings

• gotcha WebDataset implements PyTorch's `IterableDataset` and thus does not provide a `__len__` method by default. Code expecting `len(dataset)` will raise a `TypeError`. To provide a length, you must explicitly add `with_length(N)` to your pipeline. This also impacts deterministic epoch balancing in distributed training.
  Fix: Append `.with_length(N)` to your dataset pipeline if you require a length. For distributed training, consider `wds.resampled()` for approximate balancing or careful shard management.
• breaking Direct string arguments like `decode('PIL')` or `decode('numpy')` for decoding images were deprecated in favor of using actual functions (e.g., `decode(wds.decode('pil'))` or `decode('rgb')`, `decode('torchrgb')`). This change improves clarity and flexibility.
  Fix: Replace deprecated string decoders with functional calls or the appropriate shorthand string, e.g., `dataset.decode('rgb')` or `dataset.decode(wds.decode_pil)`. Consult the latest documentation for preferred decoder functions.
• gotcha Using the `pipe:` protocol with untrusted or unescaped URLs can lead to shell injection vulnerabilities, as `webdataset` executes shell commands.
  Fix: Enable secure mode by setting `webdataset.utils.enforce_security = True` in your code or by setting the environment variable `WDS_SECURE=1`. Avoid using `pipe:` with untrusted inputs.
• gotcha WebDataset relies heavily on external command-line tools like `curl`, `gsutil`, `aws`, and `file` for core I/O and type detection. This can affect portability across different operating systems or environments where these tools are not available or behave differently, and complicates error handling.
  Fix: Ensure all necessary command-line tools are installed and configured correctly in your environment. For cloud storage, consider using the `objectio` library if installed, as WebDataset passes URLs to it for direct access (without `pipe:`). Users can also implement custom `gopen_schemes`.
• gotcha Achieving precisely balanced epochs and avoiding sample repetition in multi-worker or distributed training setups (especially with `resampled=True` and shuffling) can be complex. Older usage of `repeat` argument might be outdated. Workers can endlessly repeat their shard shares without proper configuration.
  Fix: For distributed training, use `wds.split_by_node` and `wds.split_by_worker` in your pipeline. If using `resampled=True`, ensure appropriate logic to handle epoch boundaries. Be cautious with the `repeat()` method and consider its interaction with `with_epoch()` if you need fixed epoch sizes. The `wds.DataPipeline` can explicitly manage these stages.
• gotcha Long delays before the first batch, or inconsistent batch completion times, can occur due to large batch sizes, large shuffle buffers requiring time to fill, or slow underlying disk/storage access. This is often a configuration issue rather than a `webdataset` bug.
  Fix: Profile your data pipeline to identify bottlenecks. Reduce shuffle buffer size for initial debugging, ensure efficient network/disk I/O, and optimize image decoding/preprocessing steps. Monitor `curl` performance if using remote URLs.

Install

• pip install webdataset PyPI
• pip install git+https://github.com/webdataset/webdataset.git GitHub latest

Imports

• webdataset

  import webdataset as wds

• WebDataset

  dataset = wds.WebDataset(url)

  While `Dataset` might work in older versions or internal contexts, the canonical and recommended approach is to import `webdataset as wds` and use `wds.WebDataset`.

Quickstart

This quickstart demonstrates how to create a `webdataset` pipeline to load data from a remote TAR file, apply shuffling and decoding, extract specific components like images and JSON metadata, and then preprocess and batch the samples. It shows the typical 'fluid' interface with chained method calls and how it integrates with PyTorch-style data iteration. It fetches from a publicly available OpenImages shard, decodes using PIL, and extracts components into tuples.

import webdataset as wds
import torch
import os
from itertools import islice

# Example URL to a public WebDataset shard. In a real scenario, this would be your dataset path(s).
# For local files: url = "file:./my_dataset-{0000..0009}.tar"
# For cloud storage: url = "pipe:gsutil cat gs://my-bucket/dataset-{0000..0009}.tar"
url = "http://storage.googleapis.com/nvdata-openimages/openimages-train-000000.tar"

# Define a simple preprocessing function (e.g., for images and labels)
def preprocess(sample):
    # Assuming 'jpg' for image and 'json' for metadata (e.g., labels)
    image = sample['jpg']
    metadata = sample.get('json')
    
    # Example: convert image to PyTorch tensor and extract a dummy label
    # In a real scenario, you'd decode and transform the image bytes properly
    # For this example, we'll just return a placeholder tensor and label
    # (webdataset.decode() would handle actual image decoding)
    
    # If actual image decoding is not done yet, 'image' might be bytes.
    # For a quickstart without full image processing libs, mock it:
    if isinstance(image, bytes):
        # Mock a tensor, in a real app, use PIL/torchvision transforms
        processed_image = torch.randn(3, 224, 224) # e.g., C, H, W
    else:
        processed_image = image # If decode() was used earlier

    # Extract a dummy label from metadata, or just use a placeholder
    label = 0 # Placeholder label
    if metadata and isinstance(metadata, dict) and 'annotations' in metadata:
        try:
            label = metadata['annotations'][0]['category_id']
        except (IndexError, KeyError):
            pass

    return processed_image, label

# Create a WebDataset pipeline
dataset = (
    wds.WebDataset(url) # Load from URL
    .shuffle(100)        # Shuffle samples within a buffer
    .decode("pil")       # Decode images using PIL (requires Pillow installed)
    .to_tuple("jpg", "json") # Extract 'jpg' and 'json' components as a tuple
    .map(preprocess)    # Apply custom preprocessing
    .batched(16)        # Batch samples
)

# Use with PyTorch DataLoader (optional, for parallel loading and iteration)
# If you don't use PyTorch, you can iterate directly over 'dataset'
# from torch.utils.data import DataLoader
# dataloader = DataLoader(dataset, num_workers=4, batch_size=None) # batch_size=None if .batched() is used above

print(f"Accessing the first 2 batches from: {url}")

# Iterate over a few batches
for i, (images, labels) in enumerate(islice(dataset, 2)):
    print(f"Batch {i+1}:")
    print(f"  Images shape: {images.shape}")
    print(f"  Labels: {labels}")
    if i == 1:
        break

print("Quickstart complete.")