PyTorch DataLoader Optimization Guide

by

The DataLoader is often the bottleneck that prevents your GPU from training at full utilisation. If your GPU sits at 30–60% utilisation between batches, the problem is almost certainly data loading — the CPU cannot prepare and transfer batches fast enough to keep the GPU busy. This guide covers every meaningful optimisation: worker count, pinned memory, prefetching, dataset design, and the transforms pipeline.

Baseline: Measuring Actual Throughput

Before optimising, measure. This loop isolates DataLoader throughput from model compute so you know whether the bottleneck is loading or training:

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

loader = DataLoader(dataset, batch_size=256, shuffle=True, num_workers=4)

# Measure pure data loading speed
t0 = time.perf_counter()
for i, (x, y) in enumerate(loader):
    if i == 50: break
    x = x.cuda(non_blocking=True)   # include transfer in the measurement
elapsed = time.perf_counter() - t0
print(f"50 batches in {elapsed:.2f}s = {elapsed/50*1000:.1f}ms/batch")
print(f"Throughput: {50*256/elapsed:.0f} samples/sec")

Run this with different num_workers values (0, 2, 4, 8) and compare. The optimal worker count is hardware-specific — on most machines it is between 2 and 8, and more workers can actually hurt if the dataset fits in memory (CPU cache thrashing).

num_workers: The Biggest Single Lever

num_workers=0 loads data synchronously in the main process — the training loop waits for each batch to be prepared before starting the forward pass. With num_workers > 0, worker processes prepare the next batch while the GPU runs the current one. Setting this correctly is the highest-impact single change you can make to DataLoader performance:

import os
# Rule of thumb: start with num_workers = number of CPU cores / 2
num_workers = min(8, os.cpu_count() // 2)

loader = DataLoader(
    dataset,
    batch_size=256,
    shuffle=True,
    num_workers=num_workers,
    persistent_workers=True,   # keep workers alive between epochs (PyTorch 1.7+)
    prefetch_factor=2,         # batches to prefetch per worker (default: 2)
)

persistent_workers=True avoids the overhead of spawning and destroying worker processes at the start and end of each epoch — a significant saving when epochs are short or the dataset is small. Only use it with num_workers > 0.

pin_memory and non_blocking Transfers

Pinned (page-locked) memory enables asynchronous CPU-to-GPU transfers, allowing the CPU to prepare the next batch while the GPU processes the current one:

loader = DataLoader(
    dataset,
    batch_size=256,
    num_workers=4,
    pin_memory=True,       # allocate batches in pinned memory
)

# In training loop — non_blocking=True makes the transfer asynchronous
for x, y in loader:
    x = x.cuda(non_blocking=True)
    y = y.cuda(non_blocking=True)
    # GPU can start work on x while CPU prepares the next batch
    out = model(x)
    ...

pin_memory=True only helps when training on GPU. On CPU or MPS it has no effect and wastes memory. The combination of pin_memory=True and non_blocking=True together is what enables the pipelining — either alone gives only partial benefit.

Optimising the Dataset __getitem__

Every call to __getitem__ runs in a worker process. Expensive operations here — reading from disk, decoding JPEG, applying transforms — multiply by the number of samples per epoch. Common optimisations:

from PIL import Image
import torch

class FastDataset(torch.utils.data.Dataset):
    def __init__(self, paths, labels, transform=None):
        self.paths = paths
        self.labels = labels
        self.transform = transform

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

    def __getitem__(self, idx):
        # Load image — use PIL or cv2 (cv2 is often faster for large images)
        img = Image.open(self.paths[idx]).convert('RGB')

        if self.transform:
            img = self.transform(img)

        return img, self.labels[idx]

# For datasets that fit in RAM: cache everything in __init__
class CachedDataset(torch.utils.data.Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data        # load everything upfront
        self.labels = labels
        self.transform = transform

    def __getitem__(self, idx):
        x = torch.tensor(self.data[idx], dtype=torch.float32)
        if self.transform:
            x = self.transform(x)
        return x, self.labels[idx]

Moving Transforms to GPU

CPU-side transforms (random crop, flip, normalise) run in worker processes and can be a bottleneck for image datasets. PyTorch’s torchvision.transforms.v2 and the kornia library support GPU transforms that run after the batch is transferred:

import torchvision.transforms.v2 as T

# CPU transforms — applied in DataLoader workers
cpu_transform = T.Compose([
    T.RandomResizedCrop(224),
    T.RandomHorizontalFlip(),
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
])

# GPU transforms — applied after .cuda()
gpu_transform = T.Compose([
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    T.RandomErasing(p=0.1),
])

for x, y in loader:
    x = x.cuda(non_blocking=True)
    x = gpu_transform(x)       # runs on GPU — no worker bottleneck
    out = model(x)

Figure 1 — DataLoader pipeline: CPU workers → pinned memory → GPU transfer → training

CPU Workers Load + transform Pinned Memory pin_memory=True DMA Transfer non_blocking=True GPU Training model.forward()

Using WebDataset for Large-Scale Data

When your dataset is too large to fit on a local SSD, or you are training across multiple machines, reading individual files becomes a bottleneck. WebDataset solves this by packing data into tar archives that stream sequentially — much faster than random file access for spinning disks, network storage, or cloud object stores:

import webdataset as wds

dataset = (
    wds.WebDataset("s3://my-bucket/imagenet/train-{000000..001281}.tar")
    .shuffle(1000)
    .decode("pil")
    .to_tuple("jpg", "cls")
    .map_tuple(train_transform, lambda x: x)
    .batched(256)
)

loader = DataLoader(dataset, batch_size=None, num_workers=4)

Profiling DataLoader Bottlenecks

Use PyTorch’s profiler to see exactly where time is spent:

from torch.profiler import profile, record_function, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for i, (x, y) in enumerate(loader):
        if i == 20: break
        with record_function("to_gpu"):
            x = x.cuda(non_blocking=True)
        with record_function("forward"):
            out = model(x)

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
# High "DataLoader" time → increase num_workers or cache dataset
# High "to_gpu" time → enable pin_memory + non_blocking
# High "forward" time → GPU is the bottleneck, data loading is fine

Quick-Reference: DataLoader Settings

The settings that matter most, in order of impact: num_workers — set to half your CPU core count as a starting point, benchmark from there. persistent_workers=True — always enable when num_workers > 0. pin_memory=True — enable for GPU training. non_blocking=True on .cuda() calls — pairs with pin_memory. prefetch_factor — default of 2 is usually optimal; increase to 4 only if workers are consistently idle. drop_last=True — drop the last incomplete batch to avoid recompilation with variable batch sizes when using torch.compile(). Beyond DataLoader settings, caching preprocessed data to disk in a fast format (HDF5, memory-mapped NumPy arrays) and moving expensive transforms to GPU are the next highest-impact changes. If GPU utilisation is above 85% consistently, your DataLoader is no longer the bottleneck and further loading optimisation will not improve training speed.

Start by measuring GPU utilisation with nvidia-smi dmon or torch.cuda.utilization() during a training run. If it is below 70%, the DataLoader is starving the GPU — work through the optimisations above from the top. If it is above 85%, move on to model-level optimisations like torch.compile() or mixed precision. The benchmark loop at the top of this guide takes five minutes to run and immediately tells you whether DataLoader is your bottleneck and how much each change helps.

Figure 1 — DataLoader pipeline: CPU workers → pinned memory → GPU transfer → training

CPU Workers Load + transform Pinned Memory pin_memory=True DMA Transfer non_blocking=True GPU Training model.forward()

Batch Size and Gradient Accumulation

Larger batches mean fewer DataLoader iterations per epoch and less overhead — but they also require more GPU memory. Gradient accumulation lets you simulate a large effective batch size without exceeding GPU memory:

accumulation_steps = 4   # effective batch = batch_size * 4
optimizer.zero_grad()

for i, (x, y) in enumerate(loader):
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    out = model(x)
    loss = criterion(out, y) / accumulation_steps   # scale loss
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Gradient accumulation reduces the number of optimizer steps per epoch but does not reduce DataLoader overhead — the same number of batches are still loaded. For DataLoader optimisation specifically, the goal is to make each batch arrive faster, not to change how many batches there are.

Distributed Data Loading with DistributedSampler

When training on multiple GPUs with DDP, each GPU should see a disjoint subset of the data. Use DistributedSampler to partition the dataset across GPUs automatically:

from torch.utils.data.distributed import DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)

loader = DataLoader(
    dataset,
    batch_size=256,          # per-GPU batch size
    sampler=sampler,         # replaces shuffle=True
    num_workers=4,
    pin_memory=True,
    persistent_workers=True,
)

# Must set epoch at the start of each epoch for correct shuffling
for epoch in range(n_epochs):
    sampler.set_epoch(epoch)
    for x, y in loader:
        ...

With DDP, the effective batch size is batch_size × num_gpus. Adjust the learning rate accordingly — a common heuristic is to scale the LR linearly with the effective batch size (linear scaling rule), though this requires a warmup period to stabilise.

Diagnosing the Bottleneck: A Decision Tree

GPU utilisation is the key metric. Measure it with nvidia-smi dmon -s u during training and look at the SM (streaming multiprocessor) utilisation column. If utilisation is low between batches but high during the forward/backward pass, the DataLoader is the bottleneck — the GPU is idle while it waits for data. If utilisation is consistently high throughout, DataLoader is not the bottleneck. If utilisation is low throughout, the model is compute-light and you may be able to use a larger batch or more aggressive augmentation without penalty. The actions to take at each level: below 50% → increase num_workers first, then enable pin_memory and non_blocking; 50–70% → check prefetch_factor and consider GPU transforms; 70–85% → DataLoader is acceptable, optimise elsewhere; above 85% → DataLoader is not the bottleneck.

Storing Data in Fast Formats

File format choice has a large impact on load speed. Individual JPEGs from a spinning disk or network share are the worst case — each file requires a separate open and decode. Better alternatives in order of speed:

  • Memory-mapped NumPy arrays (np.memmap) — fastest for numerical data that fits on local SSD; random access with zero decode overhead.
  • HDF5 (h5py) — good for structured datasets; supports compression and chunked access; slower than memmap but more flexible.
  • LMDB — key-value store used by many large-scale vision datasets; excellent random access performance; used by ImageNet preprocessing pipelines.
  • WebDataset tar archives — best for datasets on network storage or cloud; sequential reads instead of random file I/O.
import numpy as np, h5py

# Save preprocessed dataset to HDF5 once
with h5py.File('dataset.h5', 'w') as f:
    f.create_dataset('images', data=images_array, compression='lzf')
    f.create_dataset('labels', data=labels_array)

# Fast loading in Dataset.__getitem__
class H5Dataset(torch.utils.data.Dataset):
    def __init__(self, path):
        self.f = h5py.File(path, 'r')
        self.images = self.f['images']
        self.labels = self.f['labels']
    def __len__(self): return len(self.labels)
    def __getitem__(self, idx):
        return torch.tensor(self.images[idx]), int(self.labels[idx])

Custom Collate Functions

The default collate function stacks tensors into a batch. For variable-length sequences (NLP, time series) you need a custom collate to pad sequences to the same length within each batch:

from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    sequences, labels = zip(*batch)
    # Pad to the length of the longest sequence in this batch
    padded = pad_sequence(sequences, batch_first=True, padding_value=0)
    labels = torch.tensor(labels)
    return padded, labels

loader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

Sampler Strategies

The default shuffle=True uses a random sampler that shuffles all indices at the start of each epoch. For class-imbalanced datasets, use a weighted sampler to oversample minority classes:

from torch.utils.data import WeightedRandomSampler

# Compute weight per sample (inverse class frequency)
class_counts = [500, 5000]   # class 0: 500 samples, class 1: 5000
weights = [1.0/class_counts[label] for label in all_labels]

sampler = WeightedRandomSampler(
    weights=weights,
    num_samples=len(weights),
    replacement=True
)
loader = DataLoader(dataset, batch_size=256, sampler=sampler)
# Do NOT set shuffle=True when using a sampler

Weighted sampling effectively balances the class distribution seen during training without modifying the dataset. The trade-off: each epoch may see some samples multiple times and others not at all, which affects how quickly the model converges on minority class examples.

Start by measuring GPU utilisation with nvidia-smi dmon or torch.cuda.utilization() during a real training run. Below 70% consistently means the DataLoader is the bottleneck — work through the optimisations above from the top. Above 85% means the DataLoader is not the issue and further loading optimisation will not improve speed. The benchmark loop at the top takes five minutes and gives you the actual number for your specific dataset and hardware. Run it with num_workers set to 0, 2, 4, and 8 in sequence and plot the throughput — the optimal setting is almost always obvious from the numbers, and it differs enough between machines that guessing is not worth it. On a machine with fast NVMe SSD and many CPU cores, 8 workers often saturates the DataLoader; on a machine with slow HDDs or a network filesystem, even 4 workers may not help because the bottleneck shifts to I/O bandwidth rather than CPU processing.

Leave a Comment