PyTorch DataLoader Performance: num_workers, pin_memory, Prefetching and Why Your GPU Is Waiting
PyTorch: Zero to Hero — Step 05
A fast model with a slow input pipeline is still a slow training system.
One of the most common PyTorch performance failures looks like this:
GPU utilization: 20% → 95% → 10% → 90% → 15%
The model is not necessarily slow.
The GPU may simply be waiting for the next batch.
This article is about finding out where the wait is happening.
We are going to build and instrument a DataLoader, then attack the usual bottlenecks:
num_workers
pin_memory
non_blocking transfers
prefetch_factor
persistent_workers
batch_size
slow transforms
Python overhead
shared-memory pressure
multiprocessing mistakes
The goal is not to memorize a magic configuration such as:
DataLoader(..., num_workers=8, pin_memory=True)
That configuration may be excellent on one machine and terrible on another.
The goal is to measure the pipeline and tune the actual bottleneck.
The pipeline you are really training
Most training diagrams show this:
batch → model → loss → backward → optimizer
But the real system is closer to:
flowchart TD
A[disk / network / RAM] --> B[Dataset.__getitem__]
B --> C[decode / parse / augment]
C --> D[collate batch]
D --> E[worker → main process]
E --> F[CPU tensor]
F --> G[CPU → GPU transfer]
G --> H[forward]
H --> I[backward]
I --> J[optimizer]
style G fill:#ffcc80,stroke:#333
style H fill:#81c784,stroke:#333
style I fill:#81c784,stroke:#333
style J fill:#81c784,stroke:#333
If any stage takes longer than the GPU computation, the accelerator waits.
That means an expensive GPU can spend a surprising amount of time doing this:
nothing
Start with a deliberately slow dataset
We need something that makes the problem visible.
import time
import torch
from torch.utils.data import Dataset, DataLoader
class SlowDataset(Dataset):
def __init__(self, size=10_000, feature_dim=1024, delay=0.002):
self.size = size
self.feature_dim = feature_dim
self.delay = delay
def __len__(self):
return self.size
def __getitem__(self, index):
# Simulate decoding, parsing, augmentation, disk work, etc.
time.sleep(self.delay)
x = torch.randn(self.feature_dim)
y = torch.randint(0, 10, ()).long()
return x, y
The sleep() is artificial, but the scheduling problem is real.
In production the delay may come from:
JPEG decoding
JSON parsing
PIL transforms
CSV parsing
network storage
compressed archives
audio decoding
random augmentation
Python tokenization
small-file filesystem overhead
Build a benchmark before tuning anything
Do not optimize a DataLoader by changing knobs and watching nvidia-smi casually.
Measure batch arrival time.
from time import perf_counter
def benchmark_loader(loader, max_batches=100):
start = perf_counter()
batches = 0
examples = 0
for x, y in loader:
batches += 1
examples += x.shape[0]
if batches >= max_batches:
break
elapsed = perf_counter() - start
return {
"batches": batches,
"examples": examples,
"seconds": elapsed,
"batches_per_second": batches / elapsed,
"examples_per_second": examples / elapsed,
}
Baseline:
dataset = SlowDataset()
loader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=0,
)
print(benchmark_loader(loader))
num_workers=0 means the main process performs the loading synchronously.
If __getitem__() is expensive, training cannot overlap that work with model execution.
num_workers=0 is not wrong
It is the safest debugging configuration.
It is also often correct when:
the dataset is already in RAM
__getitem__ is extremely cheap
the model itself is CPU-bound
multiprocessing overhead exceeds the work
running inside constrained containers
working on Windows while debugging worker failures
So do not start from:
More workers are always faster.
Start from:
Does parallel loading hide enough CPU work to justify multiprocessing?
Benchmark num_workers
Do this experimentally.
def make_loader(num_workers):
return DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=num_workers,
)
for workers in [0, 1, 2, 4, 8]:
loader = make_loader(workers)
result = benchmark_loader(loader)
print(
f"workers={workers:<2} "
f"batches/s={result['batches_per_second']:.2f} "
f"examples/s={result['examples_per_second']:.0f}"
)
Your curve may look like:
0 workers → slow
1 worker → faster
2 workers → faster
4 workers → best
8 workers → slightly worse
16 workers → much worse
Why can more workers make things slower?
Because workers compete for:
CPU cores
RAM
memory bandwidth
disk bandwidth
filesystem metadata
shared memory
Python process startup
cache locality
If eight workers all hit the same slow disk, you have not created eight disks.
We can plot this sweep to see the optimum at a glance:
import matplotlib.pyplot as plt
worker_counts = [0, 1, 2, 4, 8]
results = []
for w in worker_counts:
loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=w)
r = benchmark_loader(loader)
results.append(r['examples_per_second'])
plt.plot(worker_counts, results, marker='o')
plt.xlabel('Number of workers')
plt.ylabel('Examples per second')
plt.title('DataLoader throughput vs num_workers')
plt.grid(True)
plt.show()
The “knee” in the curve is your practical limit.
CPU oversubscription can destroy throughput
A very common failure is:
DataLoader workers × library threads >> available CPU cores
For example, each worker may invoke code backed by OpenMP, MKL, OpenCV or another threaded library.
Then:
8 workers × 8 CPU threads = 64 runnable threads
on an 8-core machine.
That can be dramatically slower than four workers using one or two threads each.
Inspect PyTorch CPU threads:
print(torch.get_num_threads())
Inside workers, you can deliberately constrain thread usage:
def worker_init_fn(worker_id):
torch.set_num_threads(1)
Then:
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
worker_init_fn=worker_init_fn,
)
This is not automatically the best value, but it is an important diagnostic when worker count makes performance unexpectedly worse.
Measure batch wait time separately from GPU compute
This is one of the most useful training-loop diagnostics you can add.
from time import perf_counter
def train_one_epoch(model, loader, optimizer, loss_fn, device):
model.train()
iterator = iter(loader)
total_data_time = 0.0
total_compute_time = 0.0
for step in range(len(loader)):
t0 = perf_counter()
x, y = next(iterator)
t1 = perf_counter()
x = x.to(device)
y = y.to(device)
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = loss_fn(logits, y)
loss.backward()
optimizer.step()
if device.type == "cuda":
torch.cuda.synchronize()
t2 = perf_counter()
total_data_time += t1 - t0
total_compute_time += t2 - t1
return {
"data_seconds": total_data_time,
"compute_seconds": total_compute_time,
}
If you get:
data_seconds = 30
compute_seconds = 10
then optimizing the model kernel is probably not your first problem.
Your accelerator is being starved.
A simple wait-ratio metric
def wait_ratio(data_seconds, compute_seconds):
total = data_seconds + compute_seconds
return data_seconds / total if total else 0.0
Interpretation:
0.05 → input pipeline is probably healthy
0.20 → worth investigating
0.50 → training spends half its time waiting for data
0.80 → model optimization is almost irrelevant until loading is fixed
The exact thresholds depend on your workload, but the metric forces the right question.
Why pin_memory=True exists
When training on CUDA, batches usually begin in normal CPU memory and must be copied to GPU memory.
Pinned host memory can support more efficient host-to-device transfers and can participate in asynchronous transfer patterns.
The normal pattern is:
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
pin_memory=True,
)
Then:
for x, y in loader:
x = x.to("cuda", non_blocking=True)
y = y.to("cuda", non_blocking=True)
Do not assume pin_memory=True will transform every workload.
If your real bottleneck is JPEG decoding or a network filesystem, shaving transfer overhead may barely move total epoch time.
Again: benchmark the whole pipeline.
Do not manually pin every batch without measuring
This is tempting:
x = x.pin_memory()
x = x.to("cuda", non_blocking=True)
But explicit pinning itself has cost.
The more conventional approach is to let DataLoader manage pinned batches:
loader = DataLoader(..., pin_memory=True)
and then use:
x = x.to(device, non_blocking=True)
for CUDA transfers.
non_blocking=True is not a magic async switch
This line:
x = x.to("cuda", non_blocking=True)
does not guarantee that your application now perfectly overlaps copies and GPU computation.
Whether useful overlap occurs depends on things such as:
source memory properties
CUDA stream behavior
dependency ordering
when synchronization occurs
whether subsequent kernels depend immediately on the copied tensor
But it is the standard building block for avoiding unnecessary host-side blocking during device copies.
The ideal overlap pattern can be seen as:
sequenceDiagram
participant CPU
participant GPU
CPU->>CPU: next batch (prefetched)
CPU->>GPU: async copy (non_blocking)
Note over CPU,GPU: CPU continues other work
GPU->>GPU: compute on previous batch
Note over GPU: copy completes in background
GPU->>GPU: compute on new batch
Benchmark transfer strategies
def copy_batches(loader, device, max_batches=100, non_blocking=False):
start = perf_counter()
for i, (x, y) in enumerate(loader):
x = x.to(device, non_blocking=non_blocking)
y = y.to(device, non_blocking=non_blocking)
if i + 1 >= max_batches:
break
if device.type == "cuda":
torch.cuda.synchronize()
return perf_counter() - start
Compare:
plain_loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
pin_memory=False,
)
pinned_loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
pin_memory=True,
)
print(copy_batches(plain_loader, device, non_blocking=False))
print(copy_batches(pinned_loader, device, non_blocking=True))
Do not publish one benchmark number from one laptop as a universal rule.
The useful output is the measurement method.
prefetch_factor: how far workers prepare ahead
With multiprocessing enabled, workers can prepare future batches while the current batch is being consumed.
Example:
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
prefetch_factor=2,
)
The idea is simple:
GPU computes batch N
workers prepare batches N+1, N+2, ...
If loading latency has spikes, extra prefetched work can help keep the consumer fed.
But prefetching consumes memory.
A rough mental model is:
memory pressure ≈ workers × prefetch_factor × batch memory
plus worker process state, dataset state and Python overhead.
If your batches are large, increasing prefetch_factor aggressively can create RAM or shared-memory problems long before it improves throughput.
Benchmark prefetch_factor
Only set it when num_workers > 0.
for prefetch in [1, 2, 4, 8]:
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
prefetch_factor=prefetch,
)
result = benchmark_loader(loader)
print(prefetch, result["examples_per_second"])
If:
1 → 5000 examples/s
2 → 7100
4 → 7200
8 → 6900
then 8 is not more optimized simply because it is larger.
persistent_workers=True
By default, worker processes may be created again when a new iterator is constructed for another epoch.
For many-epoch training, worker startup can be meaningful overhead.
Use:
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
persistent_workers=True,
)
This is especially useful when:
worker startup is expensive
dataset initialization in each process is expensive
epochs are short
training runs for many epochs
Benchmark multiple epochs, not just the first iterator.
def benchmark_epochs(loader, epochs=5):
start = perf_counter()
for _ in range(epochs):
for _batch in loader:
pass
return perf_counter() - start
Compare:
for persistent in [False, True]:
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
persistent_workers=persistent,
)
print(persistent, benchmark_epochs(loader))
The first epoch can lie to you
A first-epoch benchmark may include:
process startup
imports
filesystem cache misses
dataset initialization
kernel compilation
CUDA context initialization
page cache warm-up
If you are comparing settings, run warm-up iterations.
def warmup(loader, batches=20):
for i, _ in enumerate(loader):
if i + 1 >= batches:
break
Then benchmark.
Batch size changes DataLoader economics
Suppose each __getitem__() costs meaningful Python overhead.
A tiny batch size means more batches, more collation events and more loop overhead.
Compare:
for batch_size in [8, 16, 32, 64, 128, 256]:
loader = DataLoader(
dataset,
batch_size=batch_size,
num_workers=4,
)
result = benchmark_loader(loader)
print(batch_size, result["examples_per_second"])
But batch size also affects:
GPU memory
optimization dynamics
number of optimizer updates
kernel efficiency
activation memory
So DataLoader throughput is only one constraint.
Your transform may be the bottleneck
This is common in vision code:
class ImageDataset(Dataset):
def __getitem__(self, index):
image = load_image(self.paths[index])
image = expensive_python_transform(image)
return image
Measure inside __getitem__().
class TimedDataset(Dataset):
def __init__(self, wrapped):
self.wrapped = wrapped
self.total_time = 0.0
def __len__(self):
return len(self.wrapped)
def __getitem__(self, index):
start = perf_counter()
item = self.wrapped[index]
elapsed = perf_counter() - start
# Fine for debugging with num_workers=0.
# With multiprocessing, each worker owns its own dataset copy/state.
self.total_time += elapsed
return item
For worker-safe profiling, emit timing metadata with the sample:
class ProfiledDataset(Dataset):
def __getitem__(self, index):
t0 = perf_counter()
raw = self.read(index)
t1 = perf_counter()
transformed = self.transform(raw)
t2 = perf_counter()
return {
"x": transformed,
"read_time": t1 - t0,
"transform_time": t2 - t1,
}
Then inspect aggregated timings in the main process.
Python-heavy tokenization can starve training too
DataLoader problems are not only image problems.
Consider NLP preprocessing:
class TextDataset(Dataset):
def __getitem__(self, index):
text = self.rows[index]
tokens = slow_python_tokenizer(text)
return torch.tensor(tokens)
If tokenization is deterministic, ask whether it belongs in the hot training path at all.
Pre-tokenizing once may beat repeatedly doing expensive text work every epoch.
The same question applies to:
spectrogram generation
feature extraction
JSON parsing
resizing
normalization statistics
archive decompression
Move deterministic expensive work out of the epoch loop when practical.
Map-style vs iterable datasets matter
A map-style dataset implements operations such as:
__len__
__getitem__
An IterableDataset yields a stream.
from torch.utils.data import IterableDataset
class NumberStream(IterableDataset):
def __iter__(self):
for i in range(1_000_000):
yield torch.tensor(i)
With multiple workers, naive iterable datasets can duplicate data because each worker gets its own dataset instance.
You need worker-aware sharding.
Correctly shard an IterableDataset
from torch.utils.data import get_worker_info
class ShardedRange(IterableDataset):
def __init__(self, start, end):
self.start = start
self.end = end
def __iter__(self):
worker = get_worker_info()
if worker is None:
start = self.start
end = self.end
else:
total = self.end - self.start
per_worker = (total + worker.num_workers - 1) // worker.num_workers
start = self.start + worker.id * per_worker
end = min(start + per_worker, self.end)
yield from range(start, end)
Test it:
loader = DataLoader(
ShardedRange(0, 100),
num_workers=4,
batch_size=10,
)
values = []
for batch in loader:
values.extend(batch.tolist())
print(len(values))
print(len(set(values)))
For a correct single-pass stream you want those counts to agree.
Debug worker crashes with num_workers=0
A classic error is:
DataLoader worker exited unexpectedly
The first debugging move should usually be:
DataLoader(..., num_workers=0)
Why?
Because the real exception happens in the main process and you get the full traceback directly.
A worker wrapper can make errors feel mysterious when the underlying bug is simply:
bad file
invalid index
broken transform
shape mismatch
missing dependency
unexpected None
Windows: protect multiprocessing entry points
On spawn-based multiprocessing systems, top-level code can be re-imported in child processes.
Bad:
loader = DataLoader(dataset, num_workers=4)
for batch in loader:
train(batch)
Safer script structure:
def main():
dataset = MyDataset()
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
)
for batch in loader:
pass
if __name__ == "__main__":
main()
If multiprocessing works on Linux but recursively explodes or hangs on Windows, check the process entry point before blaming PyTorch.
Do not return CUDA tensors from DataLoader workers
Keep dataset workers producing CPU-side data.
Bad pattern:
class DatasetReturningCuda(Dataset):
def __getitem__(self, index):
x = load(index)
return x.cuda()
Better:
class CpuDataset(Dataset):
def __getitem__(self, index):
return load(index)
Then in the training process:
for x, y in loader:
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
This keeps device ownership and process boundaries much easier to reason about.
Shared-memory failures
Multi-worker loading moves data between processes.
On Linux/container setups you may encounter errors such as:
Unexpected bus error encountered in worker
or shared-memory allocation failures.
Before reaching for exotic fixes, reduce pressure:
loader = DataLoader(
dataset,
batch_size=32,
num_workers=2,
prefetch_factor=1,
)
You just reduced:
worker count
prefetched batches
batch memory
which are exactly the parameters driving much of the queue/shared-memory pressure.
In Docker, also inspect /dev/shm sizing.
Estimate batch memory
For a tensor:
def tensor_bytes(tensor):
return tensor.numel() * tensor.element_size()
For nested batch structures:
def object_bytes(obj):
if torch.is_tensor(obj):
return obj.numel() * obj.element_size()
if isinstance(obj, dict):
return sum(object_bytes(v) for v in obj.values())
if isinstance(obj, (tuple, list)):
return sum(object_bytes(v) for v in obj)
return 0
Inspect:
batch = next(iter(loader))
print(object_bytes(batch) / 1024**2, "MiB")
Then mentally multiply by queued/prefetched batches and worker count.
This is not a perfect process-memory estimate, but it is far better than tuning prefetch depth blindly.
Custom collate_fn can secretly dominate runtime
Suppose samples have variable-length sequences.
def collate_fn(samples):
# Maybe sorting, padding, Python loops, dictionaries, etc.
...
The collate function executes as part of batch construction.
Profile it independently.
samples = [dataset[i] for i in range(64)]
start = perf_counter()
batch = collate_fn(samples)
print("collate seconds:", perf_counter() - start)
If collate_fn spends 30 ms per batch and model compute spends 10 ms, your GPU will never be continuously busy without enough overlap.
A collate function should avoid needless Python loops
Bad pattern:
def slow_collate(samples):
output = []
for sample in samples:
output.append(sample["x"])
return torch.stack(output)
That particular example is simple, but real code often contains nested Python transforms, conversions and copying.
Whenever possible, use vectorized tensor operations after stacking.
Dataset object size matters with multiple workers
Worker processes may each hold dataset-related Python state.
This can hurt badly when the dataset object contains huge Python structures such as:
millions of file path strings
large dictionaries
metadata trees
huge nested lists
Be careful with designs like:
class HugeDataset(Dataset):
def __init__(self):
self.metadata = load_massive_python_dictionary()
If every worker materializes expensive independent state, increasing num_workers may explode memory usage.
Consider compact representations, memory mapping, databases, Arrow-like columnar formats, or loading worker-local state deliberately.
Do not reopen expensive resources for every sample
Bad:
def __getitem__(self, index):
connection = open_database()
row = connection.read(index)
connection.close()
return row
Better patterns depend on the resource, but often you want worker-local lazy initialization.
class DatabaseDataset(Dataset):
def __init__(self):
self.connection = None
def _get_connection(self):
if self.connection is None:
self.connection = open_database()
return self.connection
def __getitem__(self, index):
connection = self._get_connection()
return connection.read(index)
With worker processes, each worker can establish its own safe connection when first needed.
Whether this is valid depends on the database/client library, so understand its multiprocessing guarantees.
worker_init_fn and reproducible randomness
When using random transforms, worker-local random state matters.
import random
import numpy as np
def seed_worker(worker_id):
worker_seed = torch.initial_seed() % 2**32
random.seed(worker_seed)
np.random.seed(worker_seed)
Then:
generator = torch.Generator()
generator.manual_seed(1234)
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
worker_init_fn=seed_worker,
generator=generator,
)
This matters for performance experiments too.
If every benchmark run receives materially different augmentation behavior, comparing timings becomes noisier.
shuffle=True is usually not your performance problem
For ordinary map-style datasets, shuffling indices is cheap relative to real decode/transform work.
If turning off shuffle appears to create a huge speedup, investigate what the changed access order did to:
disk locality
network reads
cache behavior
compressed storage
remote object access
The sampler may have exposed a storage locality problem rather than being expensive itself.
Remote storage changes the tuning problem
For S3, network filesystems, blob storage or remote databases, worker tuning interacts with network concurrency.
Too few workers:
latency dominates
Too many:
connection contention
rate limits
bandwidth saturation
server throttling
The same benchmark loop still works.
Only the bottleneck changes.
Cache expensive deterministic work
Suppose preprocessing turns a raw sample into an expensive tensor.
raw -> decode -> resize -> feature transform -> tensor
If those operations are deterministic, consider caching the output.
For example:
from pathlib import Path
class CachedDataset(Dataset):
def __init__(self, paths, cache_dir):
self.paths = paths
self.cache_dir = Path(cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
def __getitem__(self, index):
cache_file = self.cache_dir / f"{index}.pt"
if cache_file.exists():
return torch.load(cache_file, weights_only=True)
x = expensive_preprocess(self.paths[index])
torch.save(x, cache_file)
return x
Production caching requires invalidation/versioning, but conceptually this can move expensive deterministic work out of repeated epochs.
Separate throughput from latency variance
Average batch time can hide stalls.
Record individual waits.
def batch_wait_times(loader, count=100):
iterator = iter(loader)
waits = []
for _ in range(count):
start = perf_counter()
try:
next(iterator)
except StopIteration:
break
waits.append(perf_counter() - start)
return waits
Then:
waits = batch_wait_times(loader)
waits_sorted = sorted(waits)
print("mean:", sum(waits) / len(waits))
print("max:", max(waits))
print("p95:", waits_sorted[int(len(waits_sorted) * 0.95)])
A pipeline with:
mean = 5 ms
p95 = 60 ms
may cause visible GPU utilization sawtoothing even though the average looks respectable.
Instrument GPU utilization properly
nvidia-smi is useful, but sampling can miss short idle periods.
Inside PyTorch, time operations carefully with CUDA events.
def time_cuda_step(fn):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
fn()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) # milliseconds
Example:
x, y = next(iter(loader))
x = x.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)
ms = time_cuda_step(lambda: model(x))
print(ms)
Use wall-clock timing for the end-to-end data pipeline and CUDA events for actual GPU work.
Do not confuse the two.
PyTorch Profiler can expose input stalls
A minimal profile:
from torch.profiler import profile, ProfilerActivity
activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
activities.append(ProfilerActivity.CUDA)
with profile(activities=activities, record_shapes=True) as prof:
for step, (x, y) in enumerate(loader):
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = loss_fn(logits, y)
loss.backward()
optimizer.step()
if step == 20:
break
print(
prof.key_averages().table(
sort_by="self_cpu_time_total",
row_limit=25,
)
)
The profiler will not magically tell you your architecture, but it can reveal whether time is dominated by:
Python/data loading
CPU transforms
copies
model kernels
synchronization
Add explicit profiler ranges
This makes traces much easier to read.
from torch.profiler import record_function
for x, y in loader:
with record_function("host_to_device"):
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
with record_function("forward"):
logits = model(x)
with record_function("loss"):
loss = loss_fn(logits, y)
with record_function("backward"):
loss.backward()
with record_function("optimizer"):
optimizer.step()
optimizer.zero_grad(set_to_none=True)
Now the timeline has labels that match your mental model.
Build a reusable DataLoader benchmark matrix
Instead of manually changing one variable repeatedly:
from itertools import product
def benchmark_matrix(dataset):
results = []
for workers, batch_size in product(
[0, 1, 2, 4, 8],
[32, 64, 128],
):
kwargs = {
"dataset": dataset,
"batch_size": batch_size,
"shuffle": True,
"num_workers": workers,
}
if workers > 0:
kwargs["prefetch_factor"] = 2
kwargs["persistent_workers"] = True
loader = DataLoader(**kwargs)
result = benchmark_loader(loader)
results.append({
"workers": workers,
"batch_size": batch_size,
**result,
})
return results
Print sorted results:
results = benchmark_matrix(dataset)
for row in sorted(
results,
key=lambda r: r["examples_per_second"],
reverse=True,
):
print(row)
This is already far more useful than copying someone else’s num_workers=8 from Stack Overflow.
Add pinning and prefetching to the matrix
def benchmark_configs(dataset):
configs = []
for workers in [0, 2, 4, 8]:
for batch_size in [32, 64, 128]:
if workers == 0:
configs.append({
"batch_size": batch_size,
"num_workers": 0,
"pin_memory": False,
})
continue
for prefetch in [1, 2, 4]:
for persistent in [False, True]:
configs.append({
"batch_size": batch_size,
"num_workers": workers,
"prefetch_factor": prefetch,
"persistent_workers": persistent,
"pin_memory": torch.cuda.is_available(),
})
results = []
for config in configs:
loader = DataLoader(dataset, shuffle=True, **config)
result = benchmark_loader(loader)
results.append({**config, **result})
return results
Do not run enormous parameter grids on production datasets casually, but this is a powerful local tuning pattern.
A practical training configuration function
Once measured, centralize the chosen settings.
def build_loader(
dataset,
*,
batch_size,
training,
num_workers,
device,
):
kwargs = {
"dataset": dataset,
"batch_size": batch_size,
"shuffle": training,
"num_workers": num_workers,
"pin_memory": device.type == "cuda",
}
if num_workers > 0:
kwargs.update({
"prefetch_factor": 2,
"persistent_workers": True,
})
return DataLoader(**kwargs)
This prevents configuration differences from leaking across training, validation and experiments.
Why is my GPU waiting? A debugging decision tree
Start here:
flowchart TD
A[GPU utilization low or sawtoothing] --> B[measure batch wait vs compute]
B --> C{Is batch wait large?}
C -- yes --> D[profile loader]
D --> E{try num_workers > 0}
E --> F{throughput improves?}
F -- yes --> G[tune workers]
F -- no --> H[profile __getitem__]
H --> I[storage/transform/collate?]
G --> J[CUDA training?]
I --> J
J --> K[pin_memory + non_blocking transfer]
K --> L[benchmark again]
C -- no --> M[profile model/GPU]
M --> N[kernel / architecture bottleneck]
style A fill:#ffccbc
style L fill:#c8e6c9
style M fill:#bbdefb
At every arrow:
measure again
Common failure: increasing num_workers forever
Symptoms:
CPU pinned at 100%
training slower
fans screaming
high context switching
random stalls
Fix:
for workers in [0, 1, 2, 4, 8]:
benchmark(...)
Find the knee in the curve.
Do not treat CPU cores as a target worker count.
Common failure: GPU starvation from synchronous loading
Symptoms:
GPU utilization pulses
large gaps between steps
num_workers=0
expensive decoding or augmentation
Test:
loader = DataLoader(
dataset,
batch_size=64,
num_workers=4,
)
Then compare end-to-end throughput.
Common failure: pin_memory=True does nothing
That can be completely normal.
If:
read + transform = 50 ms
CPU→GPU copy = 2 ms
then optimizing copy time cannot produce a dramatic training speedup.
This is why we separate stages instead of assuming a feature is broken.
Common failure: persistent_workers=True appears slower
Did you benchmark only one epoch?
Keeping workers alive mainly avoids repeated startup between iterators/epochs.
For one short benchmark, it may provide no advantage.
Compare:
benchmark_epochs(loader, epochs=10)
not only:
next(iter(loader))
Common failure: huge RAM growth with workers
Inspect:
number of workers
batch size
prefetch factor
sample size
dataset Python state
collated batch size
Then reduce pressure methodically.
DataLoader(
dataset,
batch_size=32,
num_workers=2,
prefetch_factor=1,
)
If RAM suddenly stabilizes, you have evidence about the bottleneck.
Common failure: workers hang indefinitely
Possible causes include:
library not fork-safe
thread locks
network client state
unsafe database connection reuse
worker deadlock
bad process entry point
blocking I/O with no timeout
First isolate:
num_workers=0
Then:
num_workers=1
Then increase gradually.
The smallest failing worker count tells you a great deal.
Common failure: validation is unnecessarily slow
Validation often does not need random augmentation.
Bad:
train_dataset = Dataset(transform=expensive_random_transform)
val_dataset = Dataset(transform=expensive_random_transform)
Better:
train_dataset = Dataset(transform=train_transform)
val_dataset = Dataset(transform=val_transform)
Also disable autograd during validation:
model.eval()
with torch.no_grad():
for x, y in val_loader:
...
That is not a DataLoader optimization, but end-to-end validation performance is what matters.
Common failure: timing without CUDA synchronization
Bad benchmark:
start = perf_counter()
output = model(x.cuda())
print(perf_counter() - start)
CUDA work is asynchronous with respect to the host in many cases.
For wall-clock GPU timing:
torch.cuda.synchronize()
start = perf_counter()
output = model(x)
torch.cuda.synchronize()
elapsed = perf_counter() - start
Without synchronization you can accidentally measure kernel launch overhead rather than actual compute time.
Build a batch inspector
def inspect_batch(batch, prefix="batch"):
if torch.is_tensor(batch):
print(
prefix,
"shape=", tuple(batch.shape),
"dtype=", batch.dtype,
"device=", batch.device,
"pinned=", batch.is_pinned() if batch.device.type == "cpu" else False,
)
return
if isinstance(batch, dict):
for key, value in batch.items():
inspect_batch(value, f"{prefix}.{key}")
return
if isinstance(batch, (tuple, list)):
for i, value in enumerate(batch):
inspect_batch(value, f"{prefix}[{i}]")
return
print(prefix, type(batch).__name__)
Use it:
batch = next(iter(loader))
inspect_batch(batch)
This catches surprising dtype/device/pinning behavior quickly.
Build a DataLoader configuration reporter
def describe_loader(loader):
attrs = [
"batch_size",
"num_workers",
"pin_memory",
"prefetch_factor",
"persistent_workers",
"drop_last",
]
for name in attrs:
print(f"{name:20} {getattr(loader, name, None)}")
Experiment logs should include loader configuration.
Otherwise six weeks later you may know:
experiment B was 18% faster
without knowing that experiment B quietly used four workers instead of zero.
Save performance metadata with experiments
import json
def save_loader_benchmark(path, config, result):
payload = {
"config": config,
"result": result,
"torch_version": torch.__version__,
"cuda_available": torch.cuda.is_available(),
}
with open(path, "w", encoding="utf-8") as f:
json.dump(payload, f, indent=2)
Performance without environment metadata is hard to reproduce.
Consider recording:
CPU model
core count
RAM
GPU
storage type
operating system
PyTorch version
CUDA version
worker count
batch size
A realistic optimized training skeleton
import torch
from torch.utils.data import DataLoader
def build_training_loader(dataset, batch_size, device):
workers = 4 # benchmark this on your machine
return DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=device.type == "cuda",
prefetch_factor=2,
persistent_workers=True,
)
def move_batch(batch, device):
x, y = batch
x = x.to(
device,
non_blocking=device.type == "cuda",
)
y = y.to(
device,
non_blocking=device.type == "cuda",
)
return x, y
def train_epoch(model, loader, optimizer, loss_fn, device):
model.train()
total_loss = 0.0
for batch in loader:
x, y = move_batch(batch, device)
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = loss_fn(logits, y)
loss.backward()
optimizer.step()
total_loss += loss.item()
return total_loss / len(loader)
The important comment is:
workers = 4 # benchmark this on your machine
not the number 4.
A benchmark-driven loader tuner
Here is a reusable starting point.
def tune_num_workers(
dataset,
batch_size,
worker_options=(0, 1, 2, 4, 8),
batches=100,
):
results = []
for workers in worker_options:
kwargs = {
"dataset": dataset,
"batch_size": batch_size,
"num_workers": workers,
"shuffle": False,
}
if workers > 0:
kwargs["persistent_workers"] = True
kwargs["prefetch_factor"] = 2
loader = DataLoader(**kwargs)
start = perf_counter()
examples = 0
count = 0
for batch in loader:
x = batch[0] if isinstance(batch, (tuple, list)) else batch
try:
examples += len(x)
except TypeError:
pass
count += 1
if count >= batches:
break
elapsed = perf_counter() - start
results.append({
"num_workers": workers,
"seconds": elapsed,
"batches_per_second": count / elapsed,
"examples_per_second": examples / elapsed if examples else None,
})
return results
Usage:
for result in tune_num_workers(dataset, batch_size=64):
print(result)
Now worker count is an empirical result.
The performance rule that matters most
Do not ask:
What is the best
num_workersvalue?
Ask:
What stage is preventing the accelerator from continuously receiving useful work?
That question scales beyond DataLoader.
It applies to:
model compilation
distributed training
network input
checkpoint writes
logging
CPU preprocessing
GPU transfers
kernel launch overhead
Performance engineering is bottleneck engineering.
Practical checklist: PyTorch DataLoader is slow
Run this in order.
1. Establish the baseline
num_workers=0
Measure examples/second.
2. Measure batch waiting separately
data wait vs model compute
3. Sweep worker count
0, 1, 2, 4, 8
4. Profile __getitem__
Separate:
read
decode
transform
5. Profile collate_fn
Especially for variable-length or nested data.
6. For CUDA, test
pin_memory=True
plus:
.to(device, non_blocking=True)
7. Tune prefetch depth
prefetch_factor=1, 2, 4
8. Benchmark persistent workers over multiple epochs
persistent_workers=True
9. Watch memory
Especially:
workers × prefetch × batch size
10. Investigate storage
If CPU parallelism does not help, your disk/network may already be saturated.
11. Check process/thread oversubscription
Workers spawning heavily threaded libraries can destroy performance.
12. Profile the entire step again
Because fixing one bottleneck often exposes the next.
Challenge: make the GPU stop waiting
Take a real dataset from one of your projects.
Create a baseline:
DataLoader(
dataset,
batch_size=YOUR_BATCH_SIZE,
num_workers=0,
pin_memory=False,
)
Record:
examples/second
batch wait time
GPU utilization
RAM usage
Then sweep:
num_workers
batch_size
prefetch_factor
persistent_workers
pin_memory
Do not change all parameters simultaneously at first.
Your final result should look something like:
baseline: 820 examples/s
4 workers: 2180 examples/s
+ persistent workers: 2390 examples/s
+ pinned memory: 2510 examples/s
Or perhaps:
baseline: 820 examples/s
4 workers: 790 examples/s
That second result is just as useful.
It tells you multiprocessing was not the bottleneck solution.
What we learned
The DataLoader is not just a convenience wrapper around a Dataset.
It is a concurrency and buffering boundary between data production and model consumption.
The important controls are:
batch_size
num_workers
pin_memory
prefetch_factor
persistent_workers
collate_fn
worker initialization
But the controls are not the lesson.
The lesson is:
measure producer speed
measure consumer speed
find the slower side
remove that bottleneck
measure again
That is how you stop an expensive GPU from waiting on Python.
Where the series goes next
Step 00 — What Are We Actually Doing?
Step 01 — Tensor Shapes and Broadcasting Bugs
Step 02 — Autograd Debugging
Step 03 — Build a Neural Network Without nn.Module
Step 04 — nn.Module, Parameters, Buffers and state_dict
Step 05 — DataLoader Performance and GPU Starvation
Step 06 — CNNs: Teaching PyTorch to See
Step 07 — Attention and Transformers From Scratch
Step 08 — Train Something Real
Step 09 — Performance, Compilation and Scale
Step 10 — Build a Small Language Model From Scratch
The next post moves back into model architecture, but we will keep the same programmer-first rule.
Instead of only explaining convolutions, we will attack the problems that show up when building them:
Conv2d input shape errors
channels-first vs channels-last confusion
calculating convolution output sizes
flatten-size mismatches
pooling mistakes
Then we will build a CNN that you can actually debug rather than merely copy.
