Preprocessing on the GPU

Setup#

pip install torch torchvision pillow numpy

Versions used: torch 2.5.1+cu121, torchvision 0.20.1+cu121, Pillow 11.0.0. This section needs an NVIDIA GPU. The measurements below were taken on an RTX A6000 with an Intel Core i9-13900K. Timings are hardware-dependent and vary between runs on the same machine; across four runs here the ratio was between 6.6 and 7.4 times, while the absolute millisecond figures moved by nearly a factor of two with background load. Treat the ratio as the finding and the milliseconds as an illustration.

Both pipelines, timed properly#

import io, time, statistics
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision.io import decode_jpeg

assert torch.cuda.is_available(), "this script needs an NVIDIA GPU"
dev = torch.device("cuda")
print(torch.cuda.get_device_name(0), "| torch", torch.__version__)

# 64 fake 1080p photos, held as encoded JPEG bytes, exactly like a camera sends.
rng = np.random.default_rng(0)
yy, xx = np.mgrid[0:1080, 0:1920]
base = (128 + 100 * np.sin(xx / 37) * np.cos(yy / 41)).astype(np.uint8)
BLOBS = []
for i in range(64):
    a = np.stack([np.roll(base, i * 7, 1), np.roll(base, i * 3, 0), base], -1)
    buf = io.BytesIO()
    Image.fromarray(a).save(buf, "JPEG", quality=85)
    BLOBS.append(buf.getvalue())

MEAN = torch.tensor([0.485, 0.456, 0.406], device=dev).view(1, 3, 1, 1)
STD = torch.tensor([0.229, 0.224, 0.225], device=dev).view(1, 3, 1, 1)


def cpu_pipeline():
    out = []
    for b in BLOBS:
        im = Image.open(io.BytesIO(b)).convert("RGB").resize((224, 224), Image.BILINEAR)
        out.append(torch.from_numpy(np.asarray(im).copy()))
    t = torch.stack(out).permute(0, 3, 1, 2).to(dev, non_blocking=True).float() / 255
    return (t - MEAN) / STD


def gpu_pipeline():
    raw = [torch.frombuffer(bytearray(b), dtype=torch.uint8) for b in BLOBS]
    imgs = decode_jpeg(raw, device=dev)              # nvJPEG, on the GPU
    t = F.interpolate(torch.stack(imgs).float(), size=(224, 224),
                      mode="bilinear", antialias=True, align_corners=False) / 255
    return (t - MEAN) / STD


def bench(fn, n=10):
    fn(); torch.cuda.synchronize()                   # warm up, then time properly
    ts = []
    for _ in range(n):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        fn()
        torch.cuda.synchronize()
        ts.append((time.perf_counter() - t0) * 1000)
    return statistics.median(ts)


a, b = cpu_pipeline(), gpu_pipeline()
print(f"\nboth pipelines produce {tuple(a.shape)} on {a.device}")
print(f"largest disagreement between them: {(a - b).abs().max().item():.4f}")

c_ms, g_ms = bench(cpu_pipeline), bench(gpu_pipeline)
print(f"\nCPU decode + resize for 64 images : {c_ms:7.1f} ms  "
      f"({64_000/c_ms:6.0f} images/second)")
print(f"GPU decode + resize for 64 images : {g_ms:7.1f} ms  "
      f"({64_000/g_ms:6.0f} images/second)")
print(f"speed-up: {c_ms/g_ms:.1f}x")

model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 32, 3, stride=2, padding=1), torch.nn.ReLU(),
    torch.nn.Conv2d(32, 64, 3, stride=2, padding=1), torch.nn.ReLU(),
    torch.nn.AdaptiveAvgPool2d(1), torch.nn.Flatten(), torch.nn.Linear(64, 10),
).to(dev).eval()

with torch.no_grad():
    m_ms = bench(lambda: model(a))
print(f"\nthe model forward pass itself : {m_ms:7.1f} ms")
print(f"preprocessing share of the total, CPU path: {c_ms/(c_ms+m_ms):.1%}")
print(f"preprocessing share of the total, GPU path: {g_ms/(g_ms+m_ms):.1%}")

NVIDIA RTX A6000 | torch 2.5.1+cu121

both pipelines produce (64, 3, 224, 224) on cuda:0
largest disagreement between them: 0.0637

CPU decode + resize for 64 images :   642.8 ms  (   100 images/second)
GPU decode + resize for 64 images :    95.1 ms  (   673 images/second)
speed-up: 6.8x

the model forward pass itself :     2.2 ms
preprocessing share of the total, CPU path: 99.7%
preprocessing share of the total, GPU path: 97.7%

The number that should stop you#

The model takes 2.2 milliseconds. Preprocessing takes 643. Preprocessing is 99.7 percent of the work on the CPU path, and still 97.7 percent after moving it to the GPU.

Read that again in terms of effort. A team could spend a month getting this model from 2.2 ms to 1.1 ms and improve end-to-end throughput by about 0.2 percent. The same team could spend an afternoon on the decode path and get nearly seven times.

This model is small, so the ratio is extreme. On a full ResNet-50 at batch 64 the forward pass is tens of milliseconds rather than two, and preprocessing is still a large share. The direction of the result generalises even where the size of it does not.

Moving to the GPU gave 6.8 times here. Two mechanisms, roughly equally: nvJPEG decodes on dedicated hardware in parallel across the batch, and the resize runs on thousands of cores instead of one. A third benefit does not show in this timing at all — only the compressed bytes cross the PCIe link, not the decoded 1080p frames.

And there is the disagreement: 0.0637. The two pipelines produce measurably different tensors. That is 16 grey levels out of 255, on a normalised tensor. It comes from Pillow's resampling filter differing from F.interpolate with antialias=True, and from decode differences between libjpeg-turbo and nvJPEG. Move preprocessing to the GPU and you have changed the model's input distribution. Re-run your accuracy evaluation. This is exactly the failure described in train and serve preprocessing mismatch, arriving through the back door of an optimisation.

How to time GPU code without lying to yourself#

Three rules, all visible in bench above.

Warm up first. The first call allocates memory, loads kernels and initialises the nvJPEG context. Timing it measures startup, not steady state.

Synchronise before and after. CUDA calls are asynchronous. Without torch.cuda.synchronize() you are timing how long it takes to queue the work, which is close to zero and completely meaningless.

Report the median, not the mean. One scheduling hiccup ruins a mean. Better still, report the median and the 95th percentile, because tail latency is what your users experience.

Overlap the copy with the compute#

Decoding on the GPU is one of two wins. The other is never letting the GPU idle while data moves.

stream = torch.cuda.Stream()
for batch in loader:
    with torch.cuda.stream(stream):
        gpu_batch = batch.to("cuda", non_blocking=True)   # needs pinned memory
    torch.cuda.current_stream().wait_stream(stream)
    output = model(gpu_batch)

non_blocking=True is a no-op unless the source tensor is in pinned (page-locked) memory. In PyTorch that means DataLoader(..., pin_memory=True), or .pin_memory() on the tensor. Set the flag without pinning and you get a synchronous copy plus the illusion that you optimised something.

The heavier tools#

For sustained high-throughput pipelines, NVIDIA DALI is the production answer. It builds a preprocessing graph — decode, crop, resize, normalise, augment — and executes it on the GPU with its own threading, plugging into PyTorch, TensorFlow and Triton.

pip install --extra-index-url https://pypi.nvidia.com nvidia-dali-cuda120

No output block: DALI is not installed on this machine, and a fabricated throughput table would be worse than none. What is worth knowing before you reach for it: DALI's resize and JPEG decode do not match Pillow's bit-for-bit either, so the same re-validation applies; and its main advantage over the code above is not raw speed but pipelining, prefetching and handling ragged input sizes without Python overhead.

CV-CUDA is the other option, offering GPU implementations of a wider set of OpenCV-style operations for when your preprocessing is more than decode-resize-normalise.

Common mistakes#

Optimising this before measuring it. Profile end to end first. If preprocessing is 5 percent of your time, moving it to the GPU is a week of work for a 5 percent ceiling.

Forgetting to re-validate accuracy. The tensors changed. Your numbers may have too.

Decoding on the GPU while the model is already GPU-bound. nvJPEG uses dedicated hardware on most datacentre cards, but the resize and normalise steps compete with the model for the same cores. If the GPU is already saturated, you have moved the bottleneck rather than removed it.

Using GPU decode for tiny images. The per-call overhead dominates below roughly a few hundred pixels square. Batch them, or leave them on the CPU.

Assuming every image is a JPEG. decode_jpeg with device="cuda" handles JPEG only. PNG, WebP and 16-bit TIFF fall back to the CPU, and a mixed-format input stream will quietly take the slow path for part of every batch.

Timing without synchronisation. The most common way to publish a fake speed-up.

Try it yourself#

Change quality=85 to quality=30 when creating the blobs and rerun. Smaller files mean less decode work; predict whether the CPU or the GPU path benefits more. Then replace the tiny model with torchvision.models.resnet50() and watch the preprocessing share fall from 99.7 percent to something you have to think about.

What to learn next#

• Is the GPU waiting for data? — diagnosing this bottleneck during training.
• Train and serve preprocessing mismatch — the accuracy risk this optimisation creates.
• Latency and throughput — the two numbers this trades between.