PyTorch DataLoader Pipeline

Your GPU is ready for its next batch. Worker 1 has finished, but worker 0 is still decoding an image. Can training continue? Follow the batch IDs below to see why a completed batch can still leave the training loop waiting.

The Four Components

For a map-style dataset with automatic batching, the Dataset retrieves samples, the Sampler chooses their indices, the BatchSampler groups those indices, and collate_fn combines the fetched samples into a batch. Setting num_workers > 0 moves fetching, transforms and collation into worker processes. The training loop still runs in the main process and performs the CUDA transfer itself.

The explorer is an authored schedule based on PyTorch 2.9.0's iterator implementation, not a captured run. Its three batches each contain two samples. The structure names matter because they separate requests, completed batches and batches that are ready to return:

The default in_order=True preserves the sampler's order. A faster B1 can be held while the iterator waits for B0. This is head-of-line blocking: adding workers does not remove a slow sample from the next required batch. in_order=False changes delivery behavior and can affect reproducibility and the observed data distribution. Iterable-style datasets have different sampling and sharding rules; they are covered in Production Patterns below.

1. Dataset

The Dataset class defines how to access individual samples. The critical design decision is what you store in __init__ versus what you load in __getitem__:

from torch.utils.data import Dataset
from pathlib import Path
from PIL import Image

class ImageDataset(Dataset):
    def __init__(self, root_dir, transform):
        self.paths = sorted(Path(root_dir).glob('**/*.jpg'))  # Store paths, NOT loaded images
        classes = sorted({path.parent.name for path in self.paths})
        self.class_ids = {name: index for index, name in enumerate(classes)}
        # Supply a transform that produces same-shaped tensors for collation.
        self.transform = transform

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        with Image.open(self.paths[idx]) as image:
            img = self.transform(image.convert('RGB'))
        label = self.class_ids[self.paths[idx].parent.name]
        return img, label

2. Sampler

The Sampler determines the order of indices. shuffle=True normally selects a random sampler for a map-style dataset; batch_size then groups those indices. In DDP, each rank normally creates its own loader and uses DistributedSampler to partition the index sequence:

from torch.utils.data import DataLoader, DistributedSampler

# Single GPU
loader = DataLoader(dataset, shuffle=True)  # Uses RandomSampler internally

# Multi-GPU with DDP
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
loader = DataLoader(dataset, sampler=sampler)  # Do not also set shuffle=True.

# IMPORTANT: call sampler.set_epoch(epoch) each epoch for proper shuffling
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # Without this, every epoch sees the same order
    for batch in loader:
        train(batch)

3. Workers

Worker processes have separate Python interpreters and GILs, so Python work can run concurrently across them. They still share CPU cores, memory bandwidth and storage. More workers can introduce contention or excessive native-library threads. prefetch_factor controls how much future work the iterator dispatches:

DataLoader(
    dataset,
    num_workers=4,
    prefetch_factor=2,       # Each worker pre-loads 2 batches ahead
    persistent_workers=True  # Keep worker processes and their Dataset replicas alive
)

With automatic batching, four workers and prefetch_factor=2 allow eight dispatched batch tasks. Those tasks can be waiting, loading, collating or ready; this is not a promise that eight completed tensors always sit in a queue. Raising prefetching helps absorb variations in latency, but cannot fix an average loading rate that is slower than training.

A rough payload allowance is workers × prefetch_factor × batch_bytes. For 64 images of shape 3×224×224 in float32, a batch is 36.75 MiB; eight batch slots are 294 MiB. Worker runtimes, dataset replicas, transform temporaries, the active batch and pinning copies are additional costs. Configure prefetch_factor and persistent_workers only when num_workers > 0. Keeping workers alive also keeps their Dataset state alive; later mutations in the main-process Dataset do not automatically update those replicas.

4. Collate Function

The collate_fn combines individual samples into batches. The default works for uniform tensors, but real-world data often requires custom handling — especially variable-length sequences:

import torch

def variable_length_collate(batch):
    """Handle sequences of different lengths by padding."""
    sequences, labels = zip(*batch)
    lengths = [len(s) for s in sequences]
    padded = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True)
    return padded, torch.tensor(labels), torch.tensor(lengths)

loader = DataLoader(dataset, collate_fn=variable_length_collate)

Process boundary: with num_workers > 0, both sample fetching and collate_fn run in each worker. With zero workers, both run on the main thread, as in the example above. Expensive padding can therefore be worker-side CPU work or main-thread work depending on the loader configuration. In _MapDatasetFetcher.fetch, retrieval happens before self.collate_fn(data); datasets can also implement the batch-aware __getitems__ fast path.

Worker Throughput Profiling

The best worker count depends on the workload and the machine: decode and augmentation cost, storage latency, CPU allocation, native thread pools, batch size and GPU compute all matter. In DDP, count workers across ranks on a node, rather than treating one loader's worker count as the machine total. The model above isolates a few costs to explain the tradeoff; it does not profile ImageNet, LMDB or your hardware.

For a real measurement, warm up the pipeline, then compare complete training-step throughput with a fixed workload and several worker counts. The following sketch measures end-to-end batches per second, including loading, transfer, forward/backward and optimizer work. It assumes an existing loader and train_step(batch), which performs transfer and training. The result is specific to that workload:

import time
import torch

iterator = iter(loader)
for _ in range(10):
    train_step(next(iterator))

torch.cuda.synchronize()  # CUDA workload: finish warmup before timing.
start = time.perf_counter()
steps = 0
for _ in range(100):
    try:
        batch = next(iterator)
    except StopIteration:
        break
    train_step(batch)
    steps += 1
torch.cuda.synchronize()
elapsed = time.perf_counter() - start
print(f"{steps / elapsed:.2f} measured batches/s over {steps} steps")

Use a loader with enough batches for the warmup. Exclude or report startup separately; keep batch size, cache state and compute settings comparable. torch.profiler can then show whether gaps come from next(loader), CPU work, transfer or GPU kernels. A short Python timer around a CUDA launch alone measures enqueue time, not completed GPU work.

The Serialization Tax

IPC still has costs: queue coordination, Python metadata, object traversal and tensor storage preparation. But PyTorch does not necessarily pickle every tensor payload byte into a private copy in the receiving process. Its multiprocessing support can share CPU tensor storage and send a handle. Moving previously unshared storage into shared memory can require a copy; this is not a claim of universal zero-copy loading.

# Payload arithmetic, not a transport benchmark.
batch_bytes = 32 * 3 * 224 * 224 * 4
print(batch_bytes)                  # 19,267,584 bytes
print(batch_bytes / (1024 ** 2))   # 18.375 MiB per batch

Four workers can keep several batches in flight, but they do not multiply the size of one batch by four. Plain pickle.dumps(sample) timing also misses DataLoader's collation, shared tensor transport, queue scheduling and pinning path. The distinction is documented in PyTorch's multiprocessing notes.

When multiprocessing loses, investigate the actual data representation and work per batch. Small, cheap datasets may favor zero workers; compact collated tensors may avoid large Python object graphs. NumPy arrays are not automatically shared merely because they are NumPy arrays, and IterableDataset describes an access pattern rather than a shared-memory transport guarantee.

Memory Analysis

The start method matters. With fork, child processes initially share virtual-memory pages through copy-on-write. Accessing Python object graphs can still dirty pages through reference-count updates even if your dataset code appears read-only. Compact numeric buffers often behave differently from large lists of Python objects. With spawn, workers start fresh interpreters and receive serialized startup arguments; tensor sharing and explicit shared or memory-mapped storage can change which bytes are private.

Do not assume every platform or Python release defaults to fork. Python 3.14 changed the default on supported POSIX systems to forkserver; macOS and Windows use spawn. Inspect multiprocessing.get_context().get_start_method() or choose an appropriate multiprocessing_context explicitly. The budget above separates shared backing from additional worker-owned Dataset storage. Expand Advanced assumptions to enter an estimated GiB amount per worker; that estimate is a budget input rather than a PyTorch setting. Check aggregate private/shared memory as well as /dev/shm; summing RSS can count shared pages more than once.

pin_memory and GPU Transfer

pin_memory=True asks DataLoader to put supported CPU tensor batches in page-locked host storage before returning them. With multiple workers, a dedicated main-process thread handles pinning after worker results arrive. With zero workers, pinning happens inline. Neither case automatically transfers data to CUDA.

loader = DataLoader(dataset, num_workers=4, pin_memory=True)
for images, labels in loader:
    images = images.to('cuda', non_blocking=True)
    labels = labels.to('cuda', non_blocking=True)
    train_step(images, labels)

non_blocking=True requests a transfer without the usual host-side synchronization. Pageable transfers can involve staging, and pinned memory is important for reliable asynchronous host-to-device transfer. Actual copy/compute overlap also needs suitable hardware, separate CUDA stream scheduling and correct dependencies; flags alone do not overlap operations ordered on the same stream. Keep the source storage valid and do not mutate pinned inputs before their transfer completes. Pinning consumes host resources and is not a guaranteed 2× speedup. Compare your workload and inspect its timeline; see PyTorch's pinning and non-blocking guide.

Production Patterns

When profiling identifies file-open cost, decoding or storage access as the bottleneck, changing the data layout may help more than adding workers. Per-file loading can still be appropriate; choose a different layout for a measured reason:

# LMDB — memory-mapped database, eliminates per-file overhead
import lmdb

env = lmdb.open('dataset.lmdb', readonly=True, lock=False)
with env.begin() as txn:
    data = txn.get(key)  # Memory-mapped read — no file open/close

# WebDataset — sequential tar reads, excellent for cloud storage
import webdataset as wds

dataset = (
    wds.WebDataset("s3://bucket/shard-{000..099}.tar")
    .decode("pil")
    .to_tuple("jpg", "cls")
)

# IterableDataset — streaming without random access
from torch.utils.data import IterableDataset, get_worker_info

class StreamingDataset(IterableDataset):
    def __init__(self, data_source, transform):
        self.data_source = data_source
        self.transform = transform

    def __iter__(self):
        worker_info = get_worker_info()
        # This sketch assumes each replica can enumerate the same finite source.
        worker_id = worker_info.id if worker_info else 0
        worker_count = worker_info.num_workers if worker_info else 1
        for index, sample in enumerate(self.data_source):
            if index % worker_count != worker_id:
                continue
            yield self.transform(sample)

When to use each:

• LMDB: Random access needed, single-machine training, dataset fits on local disk
• WebDataset: Cloud storage (S3/GCS), large-scale distributed training, sequential access is acceptable
• IterableDataset: Streaming from databases, message queues, or infinite data sources. The sketch partitions across workers only; distributed training also needs rank-level sharding. Prefer assigning source shards directly when each worker reading and discarding other samples would waste I/O.

Troubleshooting: Diagnosing Bottlenecks

Further Reading

• PyTorch DataLoader Docs
• WebDataset
• FFCV: Fast Forward Computer Vision
• PyTorch Performance Tuning Guide

Related concepts