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.
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.