How 3 Lines of PyTorch Code Made My Training ~5× Faster

Introduction

A few weeks ago, I was profiling one of my PyTorch training runs — a fairly standard image classification pipeline — and something didn't add up. My GPU utilization was hovering around 30–40%, which made no sense for the model I was training.

I started digging. The model itself wasn't the bottleneck. The forward and backward passes were fast. The real problem? The DataLoader pipeline was starving the GPU.

Every training iteration, the GPU would finish its computation and then sit idle, waiting for the CPU to prepare the next batch of data. The CPU was loading images, applying transforms, allocating memory, and performing a synchronous copy to the GPU — all while the GPU did absolutely nothing.

Three small configuration changes later, my training went from ~43 seconds to ~9 seconds. Nearly a 5× speedup, and I didn't touch a single layer of the model.

The Hidden Bottleneck

Here's the thing most PyTorch tutorials won't tell you: the default DataLoader configuration is not optimized for GPU training. Out of the box, this is what happens on every iteration:

Step 1: CPU loads a batch from disk
         ↓
Step 2: CPU transfers batch to GPU (synchronous, blocking)
         ↓
Step 3: GPU runs forward + backward pass
         ↓
Step 4: CPU starts loading the next batch

Notice the problem? Steps are strictly sequential. While the CPU is loading data, the GPU is idle. While the GPU is computing, the CPU is idle. Neither device is fully utilized.

On modern hardware — where your GPU can process a batch in milliseconds — this idle time becomes the dominant cost. You end up spending more time moving data than training on it.

Why This Happens

Three things are working against you in the default configuration:

Pageable memory. By default, PyTorch allocates CPU tensors in pageable memory — standard RAM that the OS can swap to disk. When you call .to(device), CUDA has to first copy data from pageable memory to a temporary pinned (page-locked) buffer, and then transfer it to the GPU. That's two copies instead of one.

Blocking transfers. The default .to(device) call is synchronous. The CPU issues the copy command and then waits — doing nothing — until the GPU confirms the transfer is complete. Only then does the next Python statement execute.

Single-threaded loading. With num_workers=0 (the default), your main training process handles data loading itself. It reads files, decodes images, applies transforms — all on a single thread, all blocking the training loop.

The result: your expensive GPU sits idle for a significant fraction of each training step, waiting for a CPU that's juggling too many responsibilities on a single thread.

The Fix

The solution is straightforward. Three parameters, three lines of change:

train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    pin_memory=True,
    num_workers=8
)

for images, labels in train_loader:
    images = images.to(device, non_blocking=True)
    labels = labels.to(device, non_blocking=True)

Let me break down each optimization.

pin_memory=True

When you set pin_memory=True, the DataLoader allocates batch tensors in pinned (page-locked) memory instead of pageable memory.

Why does this matter? Pinned memory is memory that the OS guarantees won't be paged out to disk. CUDA can DMA (Direct Memory Access) directly from pinned memory to GPU memory without an intermediate copy. This eliminates the extra copy step and significantly speeds up host-to-device transfers.

The trade-off is slightly higher CPU memory usage, since pinned memory can't be swapped. For most training workloads, this is negligible.

non_blocking=True

Passing non_blocking=True to .to(device) makes the CPU-to-GPU transfer asynchronous.

Instead of the CPU waiting for the transfer to complete, it fires off the copy command and immediately continues executing the next Python instruction. The GPU receives the data in the background. By the time your model actually needs the tensor for computation, the transfer has already finished.

This is the key to overlapping data loading with computation. The CPU can start preparing the next batch while the current batch is being transferred and processed on the GPU.

num_workers

Setting num_workers to a value greater than 0 spawns multiple background worker processes for data loading.

Instead of the main process loading one batch at a time, multiple workers load batches in parallel. While the GPU is training on batch N, workers are already reading, decoding, and transforming batches N+1, N+2, and so on.

A good starting point is setting num_workers to the number of CPU cores available. I typically use 4–8 workers, depending on the machine and the complexity of data augmentations.

Before vs After

Pipeline Parallelism

The fundamental shift is from sequential execution to pipeline parallelism:

Before (Sequential):
═══════════════════════════════════════════════
CPU  ▓▓▓▓▓░░░░░░░▓▓▓▓▓░░░░░░░▓▓▓▓▓░░░░░░░
GPU  ░░░░░▓▓▓▓▓▓▓░░░░░▓▓▓▓▓▓▓░░░░░▓▓▓▓▓▓▓
          Batch 1       Batch 2       Batch 3

After (Pipelined):
═══════════════════════════════════════════════
CPU  ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓
GPU  ░░▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓
     ▓ = working  ░ = idle

With pipeline parallelism, the CPU and GPU work simultaneously. Workers pre-fetch and transform future batches while the GPU trains on the current one. Pinned memory and non-blocking transfers ensure the handoff between CPU and GPU is as fast as possible.

My Results

I benchmarked this on a training run with 5 epochs on an image dataset:

That's roughly a 5× speedup from three parameter changes.

A caveat: The exact improvement depends on your dataset size, storage speed (SSD vs HDD), CPU core count, GPU model, data augmentation complexity, and how I/O-bound your pipeline already is. If your GPU is already the bottleneck (e.g., training a massive model on tiny data), these changes won't help much. Always benchmark your specific setup.

When This Helps Most

You'll see the biggest gains when:

• You're working with large image or video datasets that require heavy I/O
• You're training on a CUDA GPU with significant compute capability
• Your pipeline includes heavy preprocessing (resizing, normalization, augmentation)
• You have data augmentation transforms (random crops, flips, color jitter)
• Your machine has multi-core CPUs that workers can utilize
• Your GPU is fast relative to your data loading speed

When It Won't Make Much Difference

These optimizations have minimal impact if:

• You're training on CPU only (no GPU transfers to optimize)
• Your dataset is tiny and fits entirely in memory
• Your model is very large and GPU compute dominates (the GPU is already the bottleneck)
• Your pipeline is already optimized with custom data loaders or DALI

Key Takeaways

Conclusion

Sometimes the biggest performance improvements don't come from changing your neural network — they come from optimizing the pipeline feeding it.

If you're training models with PyTorch, give these three small changes a try and benchmark the results. You might be surprised by how much performance you're leaving on the table.