It is common to buy or rent a powerful GPU and then watch its utilisation hover well below what it should be. Very often the model is not the bottleneck at all. The GPU is waiting for the next batch because data loading and preprocessing on the CPU cannot keep up. Tuning the PyTorch DataLoader is one of the cheapest performance wins available.
First, confirm the bottleneck
Before changing settings, measure. Watch GPU utilisation while training; if it drops between steps, the input pipeline is suspect. A simple test is to time a loop over the DataLoader alone, with no model, and compare it with the time per training step. The PyTorch profiler can also show where each step spends its time.
import time
start = time.perf_counter()
for i, batch in enumerate(loader):
if i == 200:
break
print((time.perf_counter() - start) / 200, "s per batch")Use multiple workers
By default, data is loaded in the main process. Setting num_workers starts several worker processes that prepare batches in parallel. A reasonable starting point is a value close to the number of CPU cores available to the job, then adjust while measuring. Too many workers can cause memory pressure and contention.
loader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=8,
pin_memory=True,
persistent_workers=True,
prefetch_factor=4,
)Pin memory and transfer asynchronously
With pin_memory=True, batches are placed in page-locked host memory, which allows faster copies to the GPU. Pair it with non_blocking=True when moving tensors so the copy can overlap with computation.
images = images.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)Keep workers alive between epochs
Starting worker processes has a cost, especially when each worker opens files or builds lookup tables. persistent_workers=True keeps them running across epochs. prefetch_factor controls how many batches each worker prepares ahead of time.
Make the Dataset itself efficient
- Avoid heavy work in __init__ that gets copied to each worker. Load large indexes lazily or use memory-mapped files.
- Store data in efficient formats. Thousands of tiny files on a network drive are slow; sharded archives or memory-mapped arrays are usually faster.
- Decode and resize once. If every epoch resizes the same large images, preprocess them to the training resolution in advance.
- Move expensive augmentations to the GPU when the CPU is the bottleneck.
Choose a sensible batch size
Larger batches make better use of the GPU but need more memory and can change training dynamics. Mixed precision training reduces memory usage and often lets you increase the batch size. When you change batch size significantly, revisit the learning rate.
Watch out for common pitfalls
- Random number generators in workers need proper seeding, or augmentations may repeat across workers. Use a worker_init_fn or a seeded generator.
- Opening a database or file handle in __init__ and sharing it across workers can fail. Open it lazily in each worker.
- Python lists of many objects in the Dataset can increase memory use over time in workers; NumPy arrays are friendlier.
Measure again
Change one setting at a time and record throughput in samples per second. Gains that look obvious on paper sometimes vanish on a particular machine, and the best number of workers on a laptop differs from that on a cloud instance.
Distributed and multi-GPU training
When training across several GPUs, each process needs its own share of the data. Use a DistributedSampler so processes do not load the same samples, and call its set_epoch method at the start of each epoch so shuffling differs between epochs. Remember that worker counts now apply per process, so the total number of CPU workers can grow quickly.
Wrapping up
A fast model needs a fast feed. Confirm the bottleneck, add workers, pin memory, keep workers persistent, prefetch, and make the dataset cheap to read. These changes require no new hardware and can noticeably shorten training runs.
