Visualize how a DataLoader takes a raw dataset, shuffles it, applies on-the-fly augmentations, and groups it into mini-batches before feeding it to the Neural Network.
The Data Pipeline
1. Raw Dataset (Storage)
Apply Transforms & Collate
2. DataLoader Output (RAM/VRAM)
Processing Batch: 0 / 0
Epoch Progress0%
Overview
What is a DataLoader?
In deep learning frameworks like PyTorch or TensorFlow, a DataLoader is an utility class that handles the heavy lifting of preparing your data. Instead of writing complex loops to read pics from your hard drive one by one, the DataLoader fetches the pics, preprocesses them, optionally shuffles them. And groups them into perfectly sized chunks called mini-batches to feed into your CNN.
The problem a data loader solves is a speed mismatch. A GPU can process thousands of pics per second. But reading a JPEG from disk, decoding it, resizing it and augmenting it takes milliseconds on a single CPU core. Without prefetching, the GPU sits idle while the CPU prepares the next pic. A data loader closes that gap by with multiple worker processes to prepare batches in the background. So the next batch is ready in memory before the net finishes the current one.
A data loader has four responsibilities. It reads and decodes pics from disk. It applies transforms — resize, crop, flip, scale, turn to tensor. In the order you name. It groups transformed pics into a single tensor of shape (batch_size, channels, height, width). And it shuffles the order each epoch so that consecutive batches are not correlated. It helps gradient descent converge.
Batch size is the main param you choose. More batches use more GPU memory but give more stable gradient estimates. Less batches use less memory and add noise to the gradient. It can help escape local minima. Typical vals for pic sorting are 16, 32, 64 or 128, depending on the pic resolution and the GPU's VRAM. If training crashes with an out-of-memory error, halving the batch size is the first thing to try.
Data augmentation — random flips, rotations, colour jitter, cropping. Is applied inside the data loader at load time, not stored on disk. Each epoch the same pic is seen with other random transforms. It effectively multiplies the dataset size without duplicating files. The first pics are never changed.
Visualize how a DataLoader takes a raw dataset, shuffles it, applies on-the-fly augmentations, and groups it into mini-batches before feeding it to the Neural Network.
Why we don't load the entire dataset at once, and how dynamic augmentations make your model reliable.
What is a DataLoader?
In deep learning frameworks like PyTorch or TensorFlow, a DataLoader is an utility class that handles the heavy lifting of preparing your data. Instead of writing complex loops to read pics from your hard drive one by one, the DataLoader fetches the pics, preprocesses them, optionally shuffles them. And groups them into perfectly sized chunks called mini-batches to feed into your CNN.
The problem a data loader solves is a speed mismatch. A GPU can process thousands of pics per second. But reading a JPEG from disk, decoding it, resizing it and augmenting it takes milliseconds on a single CPU core. Without prefetching, the GPU sits idle while the CPU prepares the next pic. A data loader closes that gap by with multiple worker processes to prepare batches in the background. So the next batch is ready in memory before the net finishes the current one.
A data loader has four responsibilities. It reads and decodes pics from disk. It applies transforms — resize, crop, flip, scale, turn to tensor. In the order you name. It groups transformed pics into a single tensor of shape (batch_size, channels, height, width). And it shuffles the order each epoch so that consecutive batches are not correlated. It helps gradient descent converge.
Batch size is the main param you choose. More batches use more GPU memory but give more stable gradient estimates. Less batches use less memory and add noise to the gradient. It can help escape local minima. Typical vals for pic sorting are 16, 32, 64 or 128, depending on the pic resolution and the GPU's VRAM. If training crashes with an out-of-memory error, halving the batch size is the first thing to try.
Data augmentation — random flips, rotations, colour jitter, cropping. Is applied inside the data loader at load time, not stored on disk. Each epoch the same pic is seen with other random transforms. It effectively multiplies the dataset size without duplicating files. The first pics are never changed.
The Memory Problem (Why use batches?)
Imagine you have a dataset of 1,000,000 high-resolution pics. You cannot load all 1 million pics into your computer's RAM (or your GPU's VRAM) at the same time—it would instantly crash with an Out-of-Memory (OOM) error.
A DataLoader solves this by taking a Batch Size param. If Batch Size = 32, the DataLoader fetches 32 pics from the hard drive, sends them to the GPU, waits for the CNN to process them and update its weights. And then clears the memory to load the next 32 pics. This lets you to train on infinitely large datasets with limited hardware.
On-the-fly Data Augmentation
DataLoaders don't just load data. They are also responsible for preprocessing and augmentation. By applying random transformations (like flipping, rotating, or color shifting) to each pic as it is loaded, you artificially expand your dataset size. A single pic of a cat can be seen by the CNN as dozens of slightly other cats across multiple epochs, making the net highly reliable and stopping overfitting. The first pic on disk is never permanently altered.
Feeding the net without starving it
A GPU can process thousands of pics a second. Reading a JPEG from disk, decoding it, resizing it and augmenting it takes milliseconds on a CPU core. Without care, an costly GPU spends most of its time waiting.
A data loader is the machinery that closes that gap. Its job is to have the next batch ready in memory before the net finishes the current one.
Four responsibilities:
Read and decode pics from disk.
Transform them — resize, augment, scale, turn to tensors.
Batch them into a single tensor of shape (batch, channels, height, width).
Shuffle the order each epoch, so batches are not correlated.
from torch.utils.data import DataLoader
loader = DataLoader(
dataset,
batch_size=32,
shuffle=True, # training only
num_workers=8, # parallel CPU processes doing the loading
pin_memory=True, # faster host-to-GPU copies
prefetch_factor=2, # batches queued per worker
persistent_workers=True,
drop_last=True, # keeps every batch the same size
)
The settings that matter
num_workers is the one that often decides throughput. Zero means loading happens in the main process, in lockstep with training. The GPU idles for each batch. A reasonable starting point is the count of physical CPU cores. Or four per GPU. Too many causes memory pressure and contention.
batch_size trades memory against stability. More batches use the GPU more well and give smoother gradients. Less ones add useful noise and need less memory. It also interacts with the learning rate. The common rule is to scale the learning rate linearly with batch size.
pin_memory=True allocates page-locked memory so transfers to the GPU can be asynchronous. It is close to free and often worth it.
shuffle=True for training, False for validation. Unshuffled training data, especially if sorted by class, gives batches holding one class at a time and destroys batch normalisation's stats.
drop_last=True discards the final partial batch. This matters more than it sounds with batch normalisation. There a final batch of one or two samples gives wild stats.
Diagnosing a starved GPU
Watch GPU use with nvidia-smi in training. If it oscillates between 100% and near zero, the loader is the bottleneck, not the net.
In order of effectiveness:
Raise num_workers. The single biggest lever. And often the whole fix.
Pre-resize the pics on disk. Decoding 4000×3000 JPEGs to give 224×224 tensors wastes enormous CPU time. Resize the dataset once, offline.
Use a faster format. WebDataset, LMDB or TFRecord read sequentially and avoid millions of small-file steps. It matters especially on network storage.
Move augmentation to the GPU (Kornia, DALI) when the CPU cannot keep up.
Cache in memory if the dataset fits.
Use mixed precision — which speeds up the GPU side and so makes the loader fairly more of a bottleneck. But is worth doing anyway.
Guided Tests with This Interactive
Observe the Transforms:
Look at the Raw Dataset storage—the icons are pristine and perfectly upright. Now toggle the Random Rotation and Random Horizontal Flip checkboxes. The DataLoader Output instantly shows batches where the icons are flipped and rotated uniquely. Each epoch will generate new random variations of the first data.
Observe the Batch Grouping:
Set Dataset Size to 20 and Batch Size to 4. Look at the math: 20 / 4 = 5 iterations. Click Run Epoch. Watch how the DataLoader pulls 4 items at a time into memory. The CNN processes that single batch (highlighted in yellow), updates its weights. And moves to the next.
The Incomplete Batch:
Set Dataset Size to 22 and Batch Size to 5. The math says $22 \div 5 = 4.4$. Click Run Epoch. The DataLoader creates four full batches of 5. And one final "remainder" batch of 2. Frameworks like PyTorch have an option (drop_last=True) to discard this uneven batch if your CNN design strictly needs uniform sizes.
Turn Off Shuffling:
Uncheck Shuffle Dataset and click Run. Notice how the items are grouped perfectly sequentially (Cat, Cat... Dog, Dog... Bird). This is bad for training! If the CNN just sees Cats for the first 50 steps, it will over-optimize for Cats and unlearn what a Dog looks like.
Batch Size Trade-offs
Large Batch Size (e.g., 256): Faster training total since GPUs are highly parallelized. However, it needs a lot of VRAM and can sometimes cause the net to get stuck in "local minima" (sub-best score).Small Batch Size (e.g., 16): Uses very little VRAM. The updates to the weights are "noisy" and erratic. It really helps the net bounce out of local minima and find better generalizations.
The short of it
Epoch: One complete pass through the whole dataset.
Iteration/Step: One pass of a single batch through the CNN.
Iterations = Total Pics / Batch Size
DataLoaders stop Out-of-Memory crashes.
DataLoaders apply random augmentations on-the-fly to boost robustness.
Always shuffle your training data!
Writing a dataset
A custom dataset needs two ways: how many items there are. And how to give item i.
from torch.utils.data import Dataset
from PIL import Image
class ImageFolderDataset(Dataset):
def __init__(self, paths, labels, transform=None):
self.paths, self.labels, self.transform = paths, labels, transform
def __len__(self):
return len(self.paths)
def __getitem__(self, i):
img = Image.open(self.paths[i]).convert("RGB") # convert: some are grey
if self.transform:
img = self.transform(img)
return img, self.labels[i]
Three details worth building in from the start. turn("RGB") handles the greyscale and RGBA files that appear in almost each real dataset. Loading lazily in __getitem__ rather than in __init__ keeps memory flat. And giving a tuple of tensors lets the default collate function batch them without a custom collator.
For imbalanced data, WeightedRandomSampler draws cases in proportion to weights you supply. It oversamples rare classes at the batch level rather than duplicating files on disk.
Splits, leaks and reproducibility
The data loader is where a few test mistakes originate.
Split before augmenting, and apply training augmentation just to the training split. A validation loader should resize, centre-crop and scale. And do nothing random.
Split by group where groups exist. Multiple photographs of the same patient, product or scene must not straddle the split. Or the net spots the subject rather than the class.
Deduplicate. Near-identical pics on both sides of a split inflate the score. Perceptual hashing finds them cheaply.
Seed all for reproducibility — and note that with multiple workers each needs its own seeding via worker_init_fn. Or a few workers generate identical augmentation sequences.
import torch, numpy as np, random
def seed_worker(worker_id):
seed = torch.initial_seed() % 2**32
np.random.seed(seed); random.seed(seed)
loader = DataLoader(dataset, worker_init_fn=seed_worker, generator=torch.Generator().manual_seed(0))
Common mistakes
num_workers=0, which serialises loading with training and leaves the GPU idle.
Augmenting the validation set, making metrics noisy and irreproducible.
Shuffling the validation loader — harmless for metrics. But it makes per-sample debugging much harder.
Loading full-resolution originals when the net wants 224×224.
Forgetting .turn("RGB"), so an occasional greyscale file gives a single-channel tensor and a shape error thousands of steps into training.
Scaling before augmenting, so colour transforms operate on standardised vals and behave unpredictably.
Getting batches to the GPU without starving it
A data loader turns a folder of files into batches of tensors. And it is where a surprising share of training time goes. This measures batching, shuffling and the producer-consumer arithmetic that decides whether your costly accelerator spends its time computing or waiting.
example_01.pyNumPy
import numpy as np
N = 23
labels = np.array([i % 3 for i in range(N)])
print("A DATASET is two methods: how many items, and give me item i.")
print(" len(dataset) -> %d" % N)
print(" dataset[7] -> (image tensor, label %d)" % labels[7])
print(" everything else -- batching, shuffling, parallel loading --")
print(" is the LOADER's job, built on top of those two methods.")
print()
BATCH = 4
print("BATCHING. %d items, batch size %d:" % (N, BATCH))
order = np.arange(N)
batches = [order[i:i + BATCH] for i in range(0, N, BATCH)]
print("%-10s %8s %-26s" % ("batch", "size", "indices"))
for b, idx in enumerate(batches):
print("%-10d %8d %-26s" % (b, len(idx), str(list(idx))))
print(" %d full batches and one of %d. that last batch causes bugs:"
% (N // BATCH, N % BATCH))
print(" - BatchNorm on a batch of 1 has zero variance")
print(" - a loss dividing by hardcoded batch size is wrong")
print(" - a fixed-shape graph may refuse it")
print(" hence drop_last=True, which discards %d items per epoch." % (N % BATCH))
print(" on the VALIDATION set drop_last silently deletes %d examples"
% (N % BATCH))
print(" from every score you report.")
Output
Shuffling matters since sorted data violates the assumption SGD is built on.
example_02.pyNumPy
import numpy as np
rng = np.random.default_rng(5)
N, BATCH = 23, 4
labels = np.array([i % 3 for i in range(N)])
print("SHUFFLING. suppose the folder is sorted by class:")
sorted_labels = np.sort(labels)
print(" sorted: " + " ".join(str(v) for v in sorted_labels))
print("%-10s %-24s %s" % ("batch", "labels in it", "classes present"))
for b in range(0, min(4 * BATCH, N), BATCH):
chunk = sorted_labels[b:b + BATCH]
print("%-10d %-24s %s" % (b // BATCH, str(list(chunk)),
str(sorted(set(chunk.tolist())))))
print(" every early batch is one class -- the model oscillates.")
perm = rng.permutation(N)
shuf = labels[perm]
print(" shuffled: " + " ".join(str(v) for v in shuf))
print("%-10s %-24s %s" % ("batch", "labels in it", "classes present"))
for b in range(0, min(4 * BATCH, N), BATCH):
chunk = shuf[b:b + BATCH]
print("%-10d %-24s %s" % (b // BATCH, str(list(chunk)),
str(sorted(set(chunk.tolist())))))
print(" each batch now samples the whole distribution, which is exactly")
print(" the assumption SGD is built on.")
print()
counts = np.bincount(perm, minlength=N)
print("Shuffling = sampling WITHOUT replacement within an epoch:")
print(" items seen exactly once: %d of %d" % (int((counts == 1).sum()), N))
print(" items seen twice or more: %d" % int((counts > 1).sum()))
print(" a WeightedRandomSampler samples WITH replacement on purpose;")
print(" when it happens by accident, about %.0f%% of the data is missed." % (100.0 / np.e))
Output
Parallel workers overlap loading with GPU compute. Here is how use scales.
example_03.pyNumPy
import numpy as np
LOAD_MS, GPU_MS = 22.0, 30.0
print("WORKER PARALLELISM. loader vs GPU timing:")
print(" load + augment one batch, one worker: %.0f ms" % LOAD_MS)
print(" GPU processes one batch: %.0f ms" % GPU_MS)
print()
print("%-10s %14s %16s %14s %14s"
% ("workers", "loader ms/b", "GPU ms/b", "actual ms/b", "GPU busy"))
for w in (0, 1, 2, 4, 8):
supply = LOAD_MS / w if w else LOAD_MS + GPU_MS
if w == 0:
actual = LOAD_MS + GPU_MS
util = GPU_MS / actual
supply_show = LOAD_MS
else:
actual = max(supply, GPU_MS)
util = GPU_MS / actual
supply_show = supply
print("%-10d %14.1f %16.1f %14.1f %13.0f%%"
% (w, supply_show, GPU_MS, actual, 100 * util))
print(" with 0 workers the GPU sits idle during loading -- %.0f%% utilisation."
% (100 * GPU_MS / (LOAD_MS + GPU_MS)))
print(" by 2 workers the loader outruns the GPU. beyond that,")
print(" more workers buy nothing because the consumer is the bottleneck.")
print(" rule: add workers until the GPU is saturated, then stop.")
Output
Bigger batches do not shrink the epoch in proportion -- and they change the gradient noise.
example_04.pyNumPy
import numpy as np
GPU_MS = 30.0
EPOCH_N = 50000
print("BATCH SIZE vs EPOCH TIME. %s samples:" % "{:,}".format(EPOCH_N))
print("%-12s %14s %16s %16s %12s"
% ("batch", "batches/ep", "cost/sample", "epoch (s)", "vs batch 4"))
base = None
for bs in (4, 16, 64, 256, 1024):
per_sample = 0.60 + 0.40 * (4.0 / bs)
nb = int(np.ceil(EPOCH_N / float(bs)))
epoch = nb * bs * per_sample * (GPU_MS / 4.0) / 1000.0
base = epoch if base is None else base
print("%-12d %14s %16.3f %16.1f %11.2fx"
% (bs, "{:,}".format(nb), per_sample, epoch, epoch / base))
print(" 4 to 1024 -- a 256x increase -- buys only ~%.0f%%."
% (100 * (1 - (0.60 + 0.40 * (4.0 / 1024)) / 1.0)))
print(" bigger batches also need a LARGER learning rate to match progress.")
print()
print("CHECKLIST, in the order these bite:")
print("%-30s %s" % ("GPU util below 90%", "add workers"))
print("%-30s %s" % ("workers help no further", "GPU is the bottleneck"))
print("%-30s %s" % ("RAM climbing each epoch", "too many workers or a leak"))
print("%-30s %s" % ("val score oddly low", "check drop_last and shuffle"))
print("%-30s %s" % ("loss spiky and slow", "data is not shuffled"))
print("%-30s %s" % ("results not reproducible", "seed workers, not just numpy"))
print(" each worker gets its own random state, so seeding only the main")
print(" process is not enough.")
Output
Questions people ask
What batch size should I use? The largest that fits in memory, up to about 256 for most vision work. Then tune the learning rate accordingly.
Why is my first epoch much slower? The operating system's file cache is cold. Later epochs read from cache and speed up.
Should I shuffle validation data? No need. Deterministic order makes debugging and comparison easier.
How do I handle pics of other sizes? Resize them in the transform. Batching needs identical shapes. Var-size batches need a custom collate function and padding.
Does the loader run on the GPU? By default no. It is CPU work in split processes. DALI and Kornia move parts of it to the GPU when the CPU is the bottleneck.
What is persistent_workers for? It keeps worker processes alive between epochs, avoiding the start-up cost each time. Worth enabling when epochs are short.
Recap in one screen
The loader's job is to keep the GPU fed: read, transform, batch and shuffle, ahead of time.
num_workers is the main throughput lever. Oscillating GPU use means the loader is the bottleneck.
Pre-resize pics offline — decoding huge JPEGs to make small tensors is wasted work.
Shuffle and augment training data just. Keep validation deterministic.
Split by group, deduplicate. And seed the workers individually for reproducibility.
Recall check
0 of 3
Say the answer out loud before you reveal it — recalling it is what makes it stick, and rereading it is not.
What does this module say about “What is a DataLoader”?
What does this module say about “The Memory Problem (Why use batches?)”?
What does this module say about “On-the-fly Data Augmentation”?
Cheat sheet
Data Loaders in CNN
See how a DataLoader batches, shuffles and augments images on the fly so the GPU never sits idle waiting for the next batch.
Ashish Jangra builds and maintains VizLearn. Every module here is written and the visualisation behind it hand-built, so the numbers in a readout come from the same code that draws the picture. Corrections are genuinely welcome and get priority over everything else — if a page states something wrong, or an animation misrepresents what the algorithm does, get in touch.