TorchData

0.11.0 · active · verified Fri Apr 10

TorchData is a Python library providing composable data loading modules for PyTorch, aiming to enhance `torch.utils.data.DataLoader` and `torch.utils.data.Dataset/IterableDataset` for scalable and performant data pipelines. It focuses on new features like `StatefulDataLoader` for checkpointing and `torchdata.nodes` for flexible data processing graphs. The current version is 0.11.0. After a period of re-evaluation, development has resumed with a focus on iterative enhancements to existing PyTorch data primitives.

Warnings

breaking DataPipes and DataLoader2, which were core components of earlier TorchData versions, have been largely removed from the library starting with version 0.9.0. They were marked as deprecated in v0.8.0. Subsequent releases, including 0.11.0, do not include or maintain these solutions.
Fix: Migrate existing pipelines to `torchdata.nodes` or leverage `StatefulDataLoader` as an enhancement to `torch.utils.data.DataLoader`. If you must use DataPipes/DataLoader2, pin your dependency to `torchdata<=0.8.0`.
breaking Python 3.8 support was dropped in TorchData v0.9.0.
Fix: Upgrade your Python environment to version 3.9 or newer. The current release (0.11.0) requires Python >=3.9.
deprecated TorchData has deprecated and removed its conda builds, as PyTorch's official conda channel itself is deprecated.
Fix: Install TorchData via pip from PyPI: `pip install torchdata`.
gotcha Be aware of specific behaviors in `StatefulDataLoader` related to `num_workers=0` and initial seeding for `RandomSampler` during state loading. These can lead to unexpected iteration patterns if not handled carefully.
Fix: Refer to the official documentation for `StatefulDataLoader` regarding checkpointing and worker/seed management to ensure deterministic and correct resumption of training. Test your checkpointing logic thoroughly.

Install

pip install torchdata Install with pip

Imports

StatefulDataLoader

from torchdata.stateful_dataloader import StatefulDataLoader

nodes
```
import torchdata.nodes as nodes
```
Represents the new direction for building data pipelines with composable iterators.
DataPipes
```
from torchdata.datapipes.iter import IterableWrapper
```
DataPipes are deprecated and largely removed starting from v0.9.0. Migrate to `torchdata.nodes` or `StatefulDataLoader`.
DataLoader2
```
from torchdata.dataloader2 import DataLoader2
```
DataLoader2 is deprecated and largely removed starting from v0.9.0. Use `StatefulDataLoader` or standard `torch.utils.data.DataLoader` instead.

Quickstart

This quickstart demonstrates how to use `StatefulDataLoader`, which is a key enhancement of `torch.utils.data.DataLoader` provided by TorchData. It's a drop-in replacement that adds checkpointing capabilities.

import torch
from torch.utils.data import TensorDataset
from torchdata.stateful_dataloader import StatefulDataLoader

# Create a dummy dataset
data = torch.randn(100, 10)
labels = torch.randint(0, 2, (100,))
dataset = TensorDataset(data, labels)

# Use StatefulDataLoader as a drop-in replacement for torch.utils.data.DataLoader
batch_size = 16
dataloader = StatefulDataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=0 # For simplicity, use 0 workers
)

print(f"Number of batches: {len(dataloader)}")

# Iterate through the data
for epoch in range(2):
    print(f"\nEpoch {epoch + 1}")
    for i, (batch_data, batch_labels) in enumerate(dataloader):
        if i % 10 == 0:
            print(f"  Batch {i}: data_shape={batch_data.shape}, labels_shape={batch_labels.shape}")
        # In a real scenario, perform training steps here

# Example of saving and loading state (checkpointing)
# This is a key feature of StatefulDataLoader
state = dataloader.state_dict()
print(f"\nSaved dataloader state: {state.keys()}")

# Simulate continued training or restart
new_dataloader = StatefulDataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)
new_dataloader.load_state_dict(state)
print("Loaded dataloader state.")

# Iteration will resume from where it left off
print("Resuming iteration (should continue from saved state):")
for i, (batch_data, batch_labels) in enumerate(new_dataloader):
    if i < 3:
        print(f"  Resumed Batch {i}: data_shape={batch_data.shape}, labels_shape={batch_labels.shape}")

view raw JSON →