torch.utils.data

Active contributors: divyanshk, ramanishsingh, scotts, aelavender

Purpose

torch.utils.data is the data-loading stack: Dataset, DataLoader, Sampler, IterableDataset, plus the multi-worker spawning, pin-memory thread, and shared-memory plumbing that turn arbitrary Python iterables into a fast pipeline of GPU-ready batches.

Directory layout

Path	Contents
torch/utils/data/__init__.py	Public surface
torch/utils/data/dataset.py	Dataset, TensorDataset, Subset, ConcatDataset, ChainDataset, IterableDataset
torch/utils/data/dataloader.py	DataLoader (~1500 lines)
torch/utils/data/sampler.py	Sampler, RandomSampler, BatchSampler, WeightedRandomSampler, DistributedSampler
torch/utils/data/_utils/	Multi-worker glue, pin-memory thread, signal handling
torch/utils/data/distributed.py	DistributedSampler
torch/utils/data/datapipes/	DataPipes (composable iterable transformations)
torch/utils/data/graph.py	DataPipe graph utilities
torch/csrc/DataLoader.cpp	C++ side of multi-worker (queues, watchdog)

Key abstractions

Type	File	Purpose
Dataset	torch/utils/data/dataset.py	__len__ + __getitem__
IterableDataset	torch/utils/data/dataset.py	__iter__-based; for streaming
DataLoader	torch/utils/data/dataloader.py	The main user-facing class
Sampler	torch/utils/data/sampler.py	Picks indices
DistributedSampler	torch/utils/data/distributed.py	Sharding sampler for DDP/FSDP
_collate_fn	torch/utils/data/_utils/collate.py	Default batching

How it works

Single-worker

loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
    train_step(batch)

In single-worker mode, DataLoader is a thin iterator that:

1. Calls the sampler to get indices.
2. Calls dataset[i] for each index.
3. Stacks the results via collate_fn.

Multi-worker (num_workers > 0)

graph LR
    Main[Main process] -->|index queue| W1[Worker 1]
    Main -->|index queue| W2[Worker 2]
    W1 -->|data queue| Pin[pin-memory thread]
    W2 -->|data queue| Pin
    Pin -->|GPU-ready batches| Main

fork or spawn is used (configurable via multiprocessing_context). Each worker:

1. Receives a list of indices via a multiprocessing.Queue.
2. Calls dataset[i] for each.
3. Sends the batched result back via another queue.

A separate pin-memory thread in the main process moves batches to pinned (page-locked) memory so the next H2D copy can be async.

The C++ side at torch/csrc/DataLoader.cpp provides a watchdog that kills the workers' children if the main process dies (avoids zombie workers).

Tensor sharing across processes

When a worker returns a tensor, the tensor's storage is shared via shared memory (CPU) or cudaIpcMemHandle (CUDA). This is built on torch.multiprocessing (torch/multiprocessing/).

IterableDataset and DataPipes

For streaming data (TFRecord-style, S3, Kafka), IterableDataset defines __iter__ and the loader doesn't know the length. The torch.utils.data.datapipes subpackage adds composable transformation primitives (a successor to torchdata patterns).

Distributed sharding

DistributedSampler(dataset) shards the index set across ranks so each rank sees roughly disjoint data. Used in DDP/FSDP loops.

Common gotchas

• num_workers=0 is the default — single-process. For non-trivial workloads bump it up.
• Worker init seed. Each worker should set its own random seed; worker_init_fn is the hook.
• CUDA in workers. Don't use CUDA in __getitem__ unless multiprocessing_context="spawn" (fork can't fork CUDA contexts).
• Pin memory only helps with CUDA. No-op on CPU-only training.
• persistent_workers=True keeps workers alive across epochs — saves significant startup overhead.

Where to look

File	Purpose
torch/utils/data/dataloader.py	DataLoader class
torch/utils/data/dataset.py	Dataset variants
torch/utils/data/sampler.py	Samplers
torch/utils/data/_utils/worker.py	Worker lifecycle
torch/utils/data/_utils/pin_memory.py	Pin-memory thread
torch/csrc/DataLoader.cpp	C++ watchdog
torch/multiprocessing/reductions.py	Tensor sharing across processes

Where to read next

• Systems / Other subsystems — torch.multiprocessing.
• Features / Distributed training — DistributedSampler.