Pipeline Performance — Pin Memory, Workers, Prefetch

If your GPU is below 80% utilization, the pipeline is the bottleneck

A perfectly-fed GPU sits near 100% during training. Anything less and your data pipeline is leaving compute on the floor. The DataLoader has the knobs to fix this — most of them you set once and forget.

The settings that matter

num_workers — number of subprocesses for parallel __getitem__. Start with the number of CPU cores; tune down if memory is tight (each worker copies the dataset). 0 = main-process loading (slow but easy to debug).
pin_memory=True — allocate batches in non-pageable CPU memory. Combined with x.to(device, non_blocking=True), enables CPU↔GPU overlap.
prefetch_factor — batches each worker prefetches ahead. Default 2; raise to 4 if your GPU is finishing batches faster than they arrive.
persistent_workers=True — keep workers alive across epochs. The first epoch pays the worker-startup tax; persistent workers pay it once for the whole training run.
drop_last=True — discard the trailing partial batch. Useful for consistent batch shapes (BatchNorm stats, distributed training).

Diagnosing the bottleneck

nvidia-smi shows GPU at 30%? Pipeline is the bottleneck.
nvidia-smi shows GPU at 95%? Pipeline is fine; speedup must come from compute (compile, AMP, bigger batch).
htop shows one CPU core at 100% and others idle? num_workers is too low.
Memory pressure during multi-worker loading? Each worker copies the dataset; lower num_workers or use IterableDataset.

Code

Production-shaped DataLoader·python

import torch
from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=8,                # try matching CPU cores
    pin_memory=True,              # fast CPU→GPU
    prefetch_factor=4,            # extra prefetch headroom
    persistent_workers=True,      # don't re-spin per epoch
    drop_last=True,               # consistent batch shape
)

device = "cuda"
for x, y in loader:
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)
    # ...training...

Computing dataset normalization stats·python

import torch
from torch.utils.data import DataLoader

def compute_mean_std(loader):
    """Compute per-channel mean and std across an image dataset."""
    n_pixels = 0
    sum_  = torch.zeros(3)
    sum_sq = torch.zeros(3)
    for x, _ in loader:
        b, c, h, w = x.shape
        n_pixels += b * h * w
        sum_  += x.sum(dim=(0, 2, 3))
        sum_sq += (x ** 2).sum(dim=(0, 2, 3))
    mean = sum_ / n_pixels
    var  = sum_sq / n_pixels - mean ** 2
    std  = torch.sqrt(var)
    return mean, std

# ImageNet defaults (memorize these): mean=[.485,.456,.406] std=[.229,.224,.225]

Quick benchmark — is your loader fast enough?·python

import time
import torch
from torch.utils.data import DataLoader

def benchmark_loader(loader, n=100, device='cuda'):
    it = iter(loader)
    t0 = time.perf_counter()
    samples = 0
    for i in range(n):
        x, _ = next(it)
        x = x.to(device, non_blocking=True)
        samples += x.size(0)
    if device == 'cuda':
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0
    print(f"{samples/elapsed:,.0f} samples/sec")

# Compare configurations
for n_workers in [0, 2, 4, 8]:
    loader = DataLoader(dataset, batch_size=64, num_workers=n_workers,
                        pin_memory=True, persistent_workers=(n_workers > 0))
    benchmark_loader(loader)

Pipeline Performance — Pin Memory, Workers, Prefetch

If your GPU is below 80% utilization, the pipeline is the bottleneck

The settings that matter

Diagnosing the bottleneck

Code

External links

Exercise

Progress

Comments 0