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.
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 images from your hard drive one by one, the DataLoader fetches the images, preprocesses them, optionally shuffles them, and groups them into perfectly sized chunks called mini-batches to feed into your CNN.
The Data Pipeline
1. Raw Dataset (Storage)
Apply Transforms & Collate
2. DataLoader Output (RAM/VRAM)
Processing Batch: 0 / 0
Epoch Progress0%
Understanding Data Loaders
Why we don't load the entire dataset at once, and how dynamic augmentations make your model robust.
The Memory Problem (Why use batches?)
Imagine you have a dataset of 1,000,000 high-resolution images. You cannot load all 1 million images 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 parameter. If Batch Size = 32, the DataLoader fetches 32 images 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 images. This allows you to train on infinitely large datasets using 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 image as it is loaded, you artificially expand your dataset size. A single image of a cat can be seen by the CNN as dozens of slightly different cats across multiple epochs, making the model highly robust and preventing overfitting. The original image on disk is never permanently altered.
Feeding the model without starving it
A GPU can process thousands of images a second. Reading a JPEG from disk, decoding it, resizing it and augmenting it takes milliseconds on a CPU core. Without care, an expensive 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 model finishes the current one.
Four responsibilities:
Read and decode images from disk.
Transform them — resize, augment, normalise, convert 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 usually decides throughput. Zero means loading happens in the main process, in lockstep with training — the GPU idles for every batch. A reasonable starting point is the number of physical CPU cores, or four per GPU. Too many causes memory pressure and contention.
batch_size trades memory against stability. Larger batches use the GPU more efficiently and give smoother gradients; smaller 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 usually worth it.
shuffle=True for training, False for validation. Unshuffled training data, especially if sorted by class, produces batches containing one class at a time and destroys batch normalisation's statistics.
drop_last=True discards the final partial batch. This matters more than it sounds with batch normalisation, where a final batch of one or two samples produces wild statistics.
Diagnosing a starved GPU
Watch GPU utilisation with nvidia-smi during training. If it oscillates between 100% and near zero, the loader is the bottleneck, not the model.
In order of effectiveness:
Raise num_workers. The single biggest lever, and often the whole fix.
Pre-resize the images on disk. Decoding 4000×3000 JPEGs to produce 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 operations, which 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 relatively more of a bottleneck, but is worth doing anyway.
Guided Experiments 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. Every epoch will generate new random variations of the original 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 architecture strictly requires 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 only 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 overall because GPUs are highly parallelized. However, it requires a lot of VRAM and can sometimes cause the model to get stuck in "local minima" (sub-optimal accuracy).Small Batch Size (e.g., 16): Uses very little VRAM. The updates to the weights are "noisy" and erratic, which actually helps the model bounce out of local minima and find better generalizations.
The short of it
Epoch: One complete pass through the entire dataset.
Iteration/Step: One pass of a single batch through the CNN.
Iterations = Total Images / Batch Size
DataLoaders prevent 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 methods: how many items there are, and how to produce 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. convert("RGB") handles the greyscale and RGBA files that appear in almost every real dataset. Loading lazily in __getitem__ rather than in __init__ keeps memory flat. And returning a tuple of tensors lets the default collate function batch them without a custom collator.
For imbalanced data, WeightedRandomSampler draws examples in proportion to weights you supply, which oversamples rare classes at the batch level rather than duplicating files on disk.
Splits, leaks and reproducibility
The data loader is where several evaluation mistakes originate.
Split before augmenting, and apply training augmentation only to the training split. A validation loader should resize, centre-crop and normalise, 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 model recognises the subject rather than the class.
Deduplicate. Near-identical images on both sides of a split inflate the score. Perceptual hashing finds them cheaply.
Seed everything for reproducibility — and note that with multiple workers each needs its own seeding via worker_init_fn, or several 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 model wants 224×224.
Forgetting .convert("RGB"), so an occasional greyscale file produces a single-channel tensor and a shape error thousands of steps into training.
Normalising before augmenting, so colour transforms operate on standardised values 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 expensive accelerator spends its time computing or waiting.
example_01.pyNumPy
import numpy as np
rng = np.random.default_rng(5)
N = 23 # deliberately not a round number
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(" that is the entire interface. everything else -- batching,")
print(" shuffling, parallel loading, pinned memory -- is the LOADER's")
print(" 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 one is the source of a"
% (N // BATCH, N % BATCH))
print(" whole family of bugs:")
print(" - BatchNorm on a batch of 1 has zero variance to normalise by")
print(" - a loss that divides by a hardcoded batch size is wrong")
print(" - a fixed-shape graph or a compiled kernel may refuse it")
print(" hence drop_last=True, which discards %d items per epoch. with"
% (N % BATCH))
print(" shuffling on it is a different %d each time, so nothing is lost"
% (N % BATCH))
print(" permanently -- but on the VALIDATION set, where you must not")
print(" shuffle, drop_last silently deletes %d examples from every"
% (N % BATCH))
print(" score you report.")
print()
print("SHUFFLING. why it is not optional -- suppose the folder is sorted")
print("by class, as folders usually are:")
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 gradient from such a batch")
print(" says 'predict class 0 always', the next says 'predict class 1")
print(" always', and the model oscillates instead of converging.")
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 is now a small sample of the whole distribution,")
print(" which is exactly the assumption stochastic gradient descent is")
print(" built on: that a batch gradient is an unbiased estimate of the")
print(" full gradient. sorted batches violate it outright.")
print()
print("A SUBTLETY WORTH KNOWING: shuffling means SAMPLING WITHOUT")
print("REPLACEMENT within an epoch. check that every item appears once:")
counts = np.bincount(perm, minlength=N)
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 loader that sampled WITH replacement would show duplicates in")
print(" an epoch and miss about %.0f%% of the dataset entirely, which is"
% (100.0 / np.e))
print(" what a WeightedRandomSampler does on purpose -- and is a bug")
print(" when it happens by accident.")
print()
print("NOW THE PART THAT COSTS REAL MONEY. the loader and the GPU are a")
print("producer and a consumer, and the slow one sets the pace. the two")
print("timings below are stand-in numbers -- measure your own -- but the")
print("arithmetic they feed is exact:")
LOAD_MS, GPU_MS = 22.0, 30.0 # per batch, per worker
print(" time to load and augment one batch, one worker: %.0f ms" % LOAD_MS)
print(" time for the GPU to process 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 # same process: strictly sequential
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, loading happens in the training process, so the")
print(" GPU sits idle for every millisecond of it -- %.0f%% utilisation."
% (100 * GPU_MS / (LOAD_MS + GPU_MS)))
print(" one worker overlaps the two, and by %d workers the loader is"
% 2)
print(" faster than the GPU and utilisation is 100%. beyond that point")
print(" more workers buy NOTHING and cost memory, because the consumer")
print(" is now the bottleneck. that is the whole tuning rule: add")
print(" workers until the GPU is saturated, then stop.")
print()
print("AND THE ARITHMETIC PEOPLE GET WRONG. suppose you raise the batch")
print("size to make training faster. a bigger batch is more EFFICIENT per")
print("sample -- one large matrix multiply beats several small ones -- but")
print("that efficiency saturates. model it and look at the epoch:")
EPOCH_N = 50000
print(" %s samples, per-sample cost falling from 1.00 toward 0.60"
% "{:,}".format(EPOCH_N))
print(" as the batch grows, which is the shape real hardware shows.")
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(" going from 4 to 1024 -- a 256-fold increase -- buys about %.0f%%."
% (100 * (1 - (0.60 + 0.40 * (4.0 / 1024)) / 1.0)))
print(" most of that arrives by batch 64, and everything after it is")
print(" nearly free of benefit. the epoch does not shrink in proportion")
print(" to the batch, because the same number of SAMPLES still has to go")
print(" through the network either way -- only the overhead per sample")
print(" falls, and overhead was never the whole cost.")
print(" what a bigger batch definitely DOES change is the gradient: it")
print(" averages over more examples, so it is less noisy, so it usually")
print(" needs a LARGER learning rate to make the same progress per step.")
print(" the common rule is to scale the learning rate with the batch")
print(" size, and forgetting it is why 'I increased the batch and")
print(" accuracy got worse' is such a common report.")
print()
print("THE CHECKLIST, in the order these bite:")
print("%-30s %s" % ("GPU utilisation below 90%", "add workers"))
print("%-30s %s" % ("workers help no further", "the 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", "check the data is shuffled"))
print("%-30s %s" % ("results not reproducible", "seed the workers, not just numpy"))
print(" the last one catches almost everyone: each worker process gets")
print(" its own copy of the random state, so seeding in the main process")
print(" is not enough, and two runs with the same seed will augment")
print(" differently unless the loader seeds each worker explicitly.")
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. Subsequent 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 images of different sizes? Resize them in the transform. Batching requires identical shapes; variable-size batches need a custom collate function and padding.
Does the loader run on the GPU? By default no — it is CPU work in separate 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 utilisation means the loader is the bottleneck.
Pre-resize images offline — decoding huge JPEGs to make small tensors is wasted work.
Shuffle and augment training data only; 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”?
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 images from your hard drive one by one, the DataLoader fetches the images, preprocesses them, optionally shuffles them, and groups them into perfectly sized chunks called mini-batches to feed into your CNN.
What does this module say about “The Memory Problem (Why use batches?)”?
Imagine you have a dataset of 1,000,000 high-resolution images. You cannot load all 1 million images 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.
What does this module say about “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 image as it is loaded , you artificially expand your dataset size. A single image of a cat can be seen by the CNN as dozens of slightly different cats across multiple epochs, making the model highly robust and preventing overfitting.
Cheat sheet
Data Loaders in CNN
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.
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.