PyTorch DataLoader

The PyTorch DataLoader is the default data feeding mechanism. Understanding its internals — worker processes, memory pinning, prefetching — is essential for diagnosing training bottlenecks.

---

1. How DataLoader Works

Main process (GPU training)
  ↓ requests batch
DataLoader
  ├── Worker 0: read data → transform → collate
  ├── Worker 1: read data → transform → collate
  ├── ...
  └── Worker N: read data → transform → collate
  ↓ prefetched batch
Pin memory (optional: CPU → pinned CPU)
  ↓
GPU transfer (cudaMemcpyAsync)
  ↓
Training step

Each worker is a separate OS process (not thread — Python GIL forces multiprocessing). Workers read data from disk/network, apply transforms, and send batches to the main process via shared memory.

---

2. Key Parameters

Parameter	Default	What it does	Tuning
num_workers	0 (main process)	Number of parallel data loading processes	4-16 per GPU (test!)
prefetch_factor	2	Batches prefetched per worker	2-4 (more hides latency)
pin_memory	False	Copy data to CUDA-pinned memory	Always True for GPU training
persistent_workers	False	Keep workers alive between epochs	True (avoids respawn cost)
drop_last	False	Drop incomplete last batch	True for distributed training

num_workers Tuning

# Too few workers: GPU starved (data loading slower than training)
# Too many workers: CPU/memory contention, diminishing returns
# Sweet spot: typically 4-8 per GPU

# Empirical test:
for nw in [0, 2, 4, 8, 16, 32]:
    loader = DataLoader(dataset, num_workers=nw, pin_memory=True)
    throughput = measure_throughput(loader)
    print(f"num_workers={nw}: {throughput:.1f} samples/sec")

Rule of thumb: Start with num_workers = 4 × num_gpus. Increase until throughput plateaus or CPU becomes saturated.

pin_memory Explained

Normal (pageable) memory → GPU transfer: kernel must copy to pinned memory first, then DMA to GPU (2 copies).

Pinned memory → GPU transfer: DMA directly to GPU (1 copy, 2× faster).

# Always use for GPU training
loader = DataLoader(dataset, pin_memory=True, num_workers=8)

Cost: Pinned memory is non-swappable — too much pins can reduce available system RAM.

---

3. Common Bottlenecks

Symptom	Cause	Fix
GPU utilization < 80%	Data loading too slow	Increase num_workers, use pin_memory
CPU at 100%	Workers saturated	Pre-process data offline, simpler transforms
High RAM usage	Too many workers + large batches	Reduce num_workers or batch size
Workers keep restarting	persistent_workers=False	Set persistent_workers=True
Slow first epoch	Workers initializing	Use persistent workers + warmup
I/O wait dominant	Disk/network too slow	Move to faster storage (NVMe, streaming)

Diagnosing the Bottleneck

import torch.profiler

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA],
) as prof:
    for batch in loader:
        train_step(batch)

# Look for: large gaps between CUDA kernels = data loading bottleneck
# Look for: DataLoader.__next__ dominating CPU time = I/O bottleneck

---

4. Limits of DataLoader

Limitation	Impact	Solution
Python multiprocessing	IPC overhead, serialization cost	Use C++ data pipelines (DALI)
Local storage assumption	Doesn’t natively stream from cloud	Use StreamingDataset, WebDataset
No cross-node coordination	Each node reads independently	Streaming datasets with sharding
Memory copies	Shared memory → main process → pinned → GPU	DALI: direct GPU loading
Single-machine design	No built-in distributed data sharding	DistributedSampler or streaming

For streaming from cloud storage, see Streaming Datasets. For GPU-native loading, see GPU-Accelerated Preprocessing.

---

Self-Check

• Draw the DataLoader pipeline (workers → shared memory → pin → GPU)
• Tune num_workers for a given GPU count and measure throughput
• Explain why pin_memory=True speeds up CPU→GPU transfer
• Diagnose a data loading bottleneck from profiler output
• Identify when DataLoader isn’t enough and alternatives are needed

---

References

• PyTorch DataLoader Documentation
• PyTorch Performance Tuning Guide