Data Concepts

Training Data Pipelines

Your GPU is probably idle. The bottleneck in most training runs is the code that fetches and decodes the next batch, not the arithmetic that learns from it.

Every stage of the input pipeline has its own throughput and the whole run moves at the slowest one, so a decode step at 360 images a second leaves a 1200-a-second accelerator idle two thirds of the time

Why Does This Exist?

Somebody buys a faster GPU and the epoch time barely moves.

The reason is almost always the same, and it's measurable in one command: nvidia-smi shows utilisation bouncing between 20% and 40%. The accelerator isn't slow, it's waiting. Every millisecond it spends idle is a millisecond spent reading a file, decoding a JPEG, or waiting for a Python worker to hand over the next batch.

A training run is a pipeline, and a pipeline moves at the rate of its slowest stage. The model's arithmetic is usually not that stage. Which makes the input pipeline a first-class part of training rather than a detail — and one of the few places where an afternoon of work can double your throughput without touching the model.

What the stages are

Read bytes from storage → decode them into arrays → apply augmentation → collate into a batch → move to the device. Each has its own rate, each can be the constraint, and only the last one is on the accelerator.

Think of It Like This

One chef and a single chopping board

A restaurant hires a chef who can plate forty covers an hour. She's genuinely that fast.

The kitchen has one prep hand, and he can wash, peel and chop for about twelve covers an hour. So the restaurant does twelve covers an hour, the chef spends most of her shift standing still, and the owner's answer is to hire a faster chef.

Nothing changes. Hire three more prep hands and the covers triple, from the same chef.

Then the interesting bit: at four prep hands the boards fill faster than she can plate, and now she's the constraint again. That's the point to stop, because a fifth prep hand buys nothing.

How It Actually Works

Find the bottleneck before changing anything

Two measurements settle it. Run the model on a single batch held in memory, in a loop, and you get the pure compute rate. Then iterate the dataloader without the model and you get the pure input rate. The smaller number is your bottleneck, and the ratio tells you how much is on the table.

People skip this and start tuning num_workers by feel. Measure first.

The moves that actually help

Prefetch and overlap. Workers should be building batch $n+1$ while the device computes on batch $n$ . That's what num_workers and a prefetch depth buy you: not raw speed, but overlap, so the two stages run concurrently instead of taking turns. Without it, your epoch is the sum of both times instead of the max.

Decode less. JPEG decoding is often the single most expensive stage. Pre-resize images to the training resolution once and store them, rather than decoding a 4000px original to make a 224px crop every epoch. For tabular and text, pre-tokenise and store the token ids.

Pick a format that reads well. Millions of small files kill throughput, because per-file overhead dominates. Batch them into large sequential shards — columnar formats for tabular, and record containers like WebDataset or TFRecord for media.

Shard properly. With $N$ workers or $N$ devices, each should read a distinct slice, and the slices should be roughly equal in bytes rather than in file count. Unequal shards mean every step waits for the slowest worker, which is a stall you'll see as periodic dips rather than a constant low.

Shuffle without random reads. True global shuffling means seeking all over storage, which is slow. The standard compromise is a shuffle buffer: read sequentially, hold $k$ examples, emit a random one, refill. Approximate but sequential. Keep the buffer large relative to your shard size, because a small buffer plus ordered shards means each batch is drawn from one narrow region of the data — and a batch that's all one class does bad things to a gradient.

Watch memory. Multiple worker processes each hold their own copy of anything they touch, so a large object referenced by the dataset class gets duplicated per worker.

Determinism costs something, and is usually worth it

A reproducible run needs a seeded shuffle, a seeded augmentation stream, and a fixed worker count, because a different number of workers interleaves batches differently. That constrains your tuning: you can't reproduce a run and change num_workers at the same time. Record the count with the run.

Worked example

The model consumes 1,200 images a second. One worker decodes 90 a second.

With four workers the input stage delivers 360 a second, so the pipeline runs at $\min(1200, 360) = 360$ and the GPU sits at 30% utilisation. One epoch of 1.2M images takes 3,333 seconds.

With sixteen workers the input stage could deliver 1,440, so the pipeline runs at 1,200 — now GPU-bound at 100%, and the epoch takes 1,000 seconds.

Same hardware, same model, 3.33× faster, from a number in a constructor. And the second measurement tells you something the first can't: going to 24 workers buys nothing at all, because the accelerator is now the constraint. That's the point to stop tuning the pipeline and start looking at the model.

Show Me the Code

1import numpy as np2
3gpu_rate: float = 1200.0      # images per second the model can consume4per_worker: float = 90.0      # images per second one worker can decode5images: int = 1_200_000       # one epoch6
7def throughput(workers: int) -> float:8    """A pipeline runs at its slowest stage, never faster."""9    return float(min(gpu_rate, workers * per_worker))10
11def epoch_seconds(workers: int) -> float:12    return images / throughput(workers)13
14print(throughput(4), round(throughput(4) / gpu_rate, 2))     # -> 360.0 0.3   GPU 30% busy15print(throughput(16), round(throughput(16) / gpu_rate, 2))   # -> 1200.0 1.0  GPU saturated16print(round(epoch_seconds(4)), round(epoch_seconds(16)))     # -> 3333 100017print(round(epoch_seconds(4) / epoch_seconds(16), 2))        # -> 3.33

min is the whole model of a pipeline. Every optimisation is either raising the smallest term or overlapping two terms so they stop adding.

Watch Out For

Tuning worker count without measuring either stage

num_workers=8 because a blog post said so. Then 16, which is worse. Then back to 4.

Without the two isolated measurements you're searching blind, and the search has real traps. Too many workers oversubscribes your CPU cores, so workers contend and throughput falls. Each worker also copies whatever the dataset object holds, so 32 workers holding a 1GB index is 32GB of RAM. And on some setups worker startup happens every epoch, so a short epoch spends a meaningful fraction of its time spawning processes.

Measure compute-only and input-only, then raise workers until input throughput exceeds compute throughput, and stop. If input throughput plateaus below compute no matter the worker count, the constraint is storage or decode, and more processes cannot fix it — that's when you pre-resize, pre-tokenise, or change the file format.

A shuffle buffer too small for how the data is laid out

Data written in class order — all the cats, then all the dogs — read sequentially with a 1,000-example shuffle buffer, on shards of 100,000.

Every batch is drawn from a window that sits entirely inside one class. So the gradient for the first stretch of the epoch says "predict cat", then abruptly "predict dog". Loss is jumpy, batch-norm statistics swing, and the run converges worse than the same data properly shuffled — with no error and nothing in the code that looks wrong.

It's the same problem as a stratified split, one level down: a batch is a sample, and a non-representative sample gives you a biased gradient.

Two fixes, and do both. Shuffle once when you write the shards, so sequential reads are already mixed. Then keep the runtime buffer large enough to span shard boundaries. Verify rather than assume: log the class distribution of the first twenty batches and check it looks like the dataset.

The Quick Version

Most training runs are input-bound, not compute-bound. Check accelerator utilisation before buying anything.
Stages: read, decode, augment, collate, transfer. The pipeline runs at the slowest one.
Measure compute-only and input-only separately. Everything else is guessing.
Workers buy overlap: batch $n+1$ is built while batch $n$ computes, so the times take a max instead of a sum.
Decode less. Pre-resize images, pre-tokenise text, and store the result.
Millions of small files are slow. Use large sequential shards, roughly equal in bytes.
A shuffle buffer approximates a global shuffle with sequential reads. Shuffle at write time too, or batches come out class-ordered.
Determinism needs a seeded shuffle, seeded augmentation and a fixed worker count. Record the count with the run.
Four workers gave 30% utilisation and a 3,333-second epoch. Sixteen gave 100% and 1,000 seconds. Twenty-four gives nothing more.