Iterable Streaming Webdataset for PyTorch from boto3 compliant storage
Reason this release was yanked:
core bugs
Project description
streaming-wds (Streaming WebDataset)
streaming-wds
is a Python library that enables efficient streaming of WebDataset-format datasets from boto3-compliant object stores for PyTorch. It's designed to handle large-scale datasets with ease, especially in distributed training contexts.
Features
- Streaming of WebDataset-format data from S3-compatible object stores
- Efficient sharding of data across both torch distributed workers and dataloader multiprocessing workers
- Supports mid-epoch resumption when used with
StreamingDataLoader
- Blazing fast data loading with local caching and explicit control over memory consumption
- Customizable decoding of dataset elements via
StreamingDataset.process_sample
Installation
You can install streaming-wds
using pip:
pip install streaming-wds
Quick Start
Here's a basic example of how to use streaming-wds:
from streaming_wds import StreamingWebDataset, StreamingDataLoader
# Create the dataset
dataset = StreamingWebDataset(
remote="s3://your-bucket/your-dataset",
split="train",
profile="your_aws_profile",
prefetch=2,
shuffle=True,
max_workers=4,
schema={".jpg": "PIL", ".json": "json"}
)
# or use a custom processing function
import torchvision.transforms.v2 as T
class ImageNetWebDataset(StreamingWebDataset):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.transforms = T.Compose([
T.ToImage(),
T.Resize((64,)),
T.ToDtype(torch.float32),
T.Normalize(mean=(128,), std=(128,)),
])
def process_sample(self, sample):
sample[".jpg"] = self.transforms(sample[".jpg"])
return sample
# Create a StreamingDataLoader for mid-epoch resumption
dataloader = StreamingDataLoader(dataset, batch_size=32, num_workers=4)
# Iterate through the data
for batch in dataloader:
# Your training loop here
pass
# You can save the state for resumption
state_dict = dataloader.state_dict()
# Later, you can resume from this state
dataloader.load_state_dict(state_dict)
Configuration
remote
: The S3 URI of your datasetsplit
: The dataset split (e.g., "train", "val", "test")profile
: The AWS profile to use for authenticationbuffer_size
: The size of the buffer for each workershuffle
: Whether to shuffle the datamax_workers
: Maximum number of worker threads for download and extractionschema
: A dictionary defining the decoding method for each data fieldcache_limit_bytes
: The maximum size of the file cache in bytes. This should be fairly large to prevent frequent cache evictions.process_sample
: Implement this function to customize the processing of each sample (for example: torchvision transforms)
Contributing
Contributions to streaming-wds are welcome! Please feel free to submit a Pull Request.
License
MIT License
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
streaming_wds-0.1.11.tar.gz
(24.1 kB
view hashes)
Built Distribution
Close
Hashes for streaming_wds-0.1.11-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | a982a5a65e69dec8e94af195db3871c38f0af05d5af5b4088ed3c5994971080a |
|
MD5 | c1a3136b3aab24c8fc879240486d14a2 |
|
BLAKE2b-256 | ef45c5c5baf9afeb7ece785084050e71f1b5ffafa3ae380a46c095433788dbe7 |