Streaming
A Data Streaming Library for Efficient Neural Network Training
Install / Use
/learn @mosaicml/StreamingREADME
👋 Welcome
We built StreamingDataset to make training on large datasets from cloud storage as fast, cheap, and scalable as possible.
It’s specially designed for multi-node, distributed training for large models—maximizing correctness guarantees, performance, and ease of use. Now, you can efficiently train anywhere, independent of your training data location. Just stream in the data you need, when you need it. To learn more about why we built StreamingDataset, read our announcement blog.
StreamingDataset is compatible with any data type, including images, text, video, and multimodal data.
With support for major cloud storage providers (AWS, OCI, GCS, Azure, Databricks, and any S3 compatible object store such as Cloudflare R2, Coreweave, Backblaze b2, etc. ) and designed as a drop-in replacement for your PyTorch IterableDataset class, StreamingDataset seamlessly integrates into your existing training workflows.
🚀 Getting Started
💾 Installation
Streaming can be installed with pip:
pip install mosaicml-streaming
🏁 Quick Start
1. Prepare Your Data
Convert your raw dataset into one of our supported streaming formats:
- MDS (Mosaic Data Shard) format which can encode and decode any Python object
- CSV / TSV
- JSONL
import numpy as np
from PIL import Image
from streaming import MDSWriter
# Local or remote directory in which to store the compressed output files
data_dir = 'path-to-dataset'
# A dictionary mapping input fields to their data types
columns = {
'image': 'jpeg',
'class': 'int'
}
# Shard compression, if any
compression = 'zstd'
# Save the samples as shards using MDSWriter
with MDSWriter(out=data_dir, columns=columns, compression=compression) as out:
for i in range(10000):
sample = {
'image': Image.fromarray(np.random.randint(0, 256, (32, 32, 3), np.uint8)),
'class': np.random.randint(10),
}
out.write(sample)
2. Upload Your Data to Cloud Storage
Upload your streaming dataset to the cloud storage of your choice (AWS, OCI, or GCP). Below is one example of uploading a directory to an S3 bucket using the AWS CLI.
<!--pytest.mark.skip-->$ aws s3 cp --recursive path-to-dataset s3://my-bucket/path-to-dataset
3. Build a StreamingDataset and DataLoader
<!--pytest.mark.skip-->from torch.utils.data import DataLoader
from streaming import StreamingDataset
# Remote path where full dataset is persistently stored
remote = 's3://my-bucket/path-to-dataset'
# Local working dir where dataset is cached during operation
local = '/tmp/path-to-dataset'
# Create streaming dataset
dataset = StreamingDataset(local=local, remote=remote, shuffle=True)
# Let's see what is in sample #1337...
sample = dataset[1337]
img = sample['image']
cls = sample['class']
# Create PyTorch DataLoader
dataloader = DataLoader(dataset)
📚 What next?
Getting started guides, examples, API references, and other useful information can be found in our docs.
We have end-to-end tutorials for training a model on:
We also have starter code for the following popular datasets, which can be found in the streaming directory:
| Dataset | Task | Read | Write | | --- | --- | --- | --- | | LAION-400M | Text and image | Read | Write | | WebVid | Text and video | Read | Write | | C4 | Text | Read | Write | | EnWiki | Text | Read | Write | | Pile | Text | Read | Write | ADE20K | Image segmentation | Read | Write | CIFAR10 | Image classification | Read | Write | | COCO | Image classification | Read | Write | | ImageNet | Image classification | Read | Write |
To start training on these datasets:
- Convert raw data into .mds format using the corresponding script from the
convertdirectory.
For example:
<!--pytest.mark.skip-->$ python -m streaming.multimodal.convert.webvid --in <CSV file> --out <MDS output directory>
- Import dataset class to start training the model.
from streaming.multimodal import StreamingInsideWebVid
dataset = StreamingInsideWebVid(local=local, remote=remote, shuffle=True)
🔑 Key Features
Seamless data mixing
Easily experiment with dataset mixtures with Stream. Dataset sampling can be controlled in relative (proportion) or absolute (repeat or samples terms). During streaming, the different datasets are streamed, shuffled, and mixed seamlessly just-in-time.
# mix C4, github code, a
Related Skills
proje
Interactive vocabulary learning platform with smart flashcards and spaced repetition for effective language acquisition.
API
A learning and reflection platform designed to cultivate clarity, resilience, and antifragile thinking in an uncertain world.
openclaw-plugin-loom
Loom Learning Graph Skill This skill guides agents on how to use the Loom plugin to build and expand a learning graph over time. Purpose - Help users navigate learning paths (e.g., Nix, German)
best-practices-researcher
The most comprehensive Claude Code skills registry | Web Search: https://skills-registry-web.vercel.app
