Learn AI Create the Future — with Pranay
Main site

Datasets and DataLoaders

What DataLoader actually does

DataLoader is four small machines in a row — a sampler choosing indices, a batcher grouping them, a fetcher calling your Dataset, and a collator gluing samples into one tensor batch.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

A DataLoader takes single samples from your Dataset and delivers them as shuffled, ready-to-train batches.

Think of a canteen at lunch hour. The kitchen — your Dataset from the last lesson — can produce any single dish on request. But nobody serves one dish at a time to three hundred people.

Instead there is a system. Someone decides today's serving order. Someone groups orders into trays of thirty. Runners fetch the dishes, and an assembler arranges each tray the same way. The kitchen keeps cooking single dishes; the system produces trays.

The DataLoader is that system, and a tray is a batch — a group of samples processed together.

Why it exists

Training has serving requirements that have nothing to do with your storage.

Models train on batches, because hardware is far faster on thirty samples at once than on thirty singles. Order must be shuffled — rearranged randomly each pass — because a model fed all ripe mangoes first and unripe last learns the arrangement instead of the fruit. And ideally the next tray gets prepared while the current one is being eaten.

None of this depends on whether your data is photos or sensor logs. So it was built once, properly, as DataLoader, and every PyTorch project on earth reuses it.

How it works

 sampler:   "today's order: 3, 41, 7, 90, ..."     (shuffled each pass)
     |
 batcher:   [3, 41, 7, 90] ... [12, 55, 2, 88]     (groups of batch_size)
     |
 fetcher:   dataset[3], dataset[41], ...            (your __getitem__, called)
     |
 collator:  four samples  →  one stacked batch      (trays, arranged identically)
     |
 your training loop:  for x, y in loader:

Four small machines in a row. Each is replaceable — later lessons swap the collator and the sampler.

A real example you have seen

Any time you have read "trained for 90 passes over 1.2 million images", this machinery — or its cousin in another framework — was the thing making that sentence physically possible.

Remember this

  • DataLoader = sampler → batcher → fetcher → collator.
  • It delivers shuffled batches; your Dataset keeps serving singles.
  • Every stage is swappable, which is where the next lessons go.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Written and tested against torch 2.5 on CPU.

The pipeline, rebuilt by hand

The fastest way to understand DataLoader is to do its job yourself once:

loader_by_hand.py
import torch
from torch.utils.data import Dataset, DataLoader, default_collate

class MangoDataset(Dataset):
    """160 fake mangoes: two measurements each, and a ripe/unripe label."""
    def __init__(self):
        g = torch.Generator().manual_seed(0)
        self.features = torch.rand(160, 2, generator=g)
        self.labels = (self.features.sum(dim=1) > 1.0).long()

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

ds = MangoDataset()

# What DataLoader does, by hand: pick indices, fetch samples, glue them.
indices = [3, 41, 7, 90]
samples = [ds[i] for i in indices]
batch_x, batch_y = default_collate(samples)
print("hand-made batch:", tuple(batch_x.shape), tuple(batch_y.shape))

loader = DataLoader(ds, batch_size=32, shuffle=True,
                    generator=torch.Generator().manual_seed(0), drop_last=False)
for i, (x, y) in enumerate(loader):
    print(f"batch {i}: x {tuple(x.shape)}  y {tuple(y.shape)}")
Output
hand-made batch: (4, 2) (4,)
batch 0: x (32, 2)  y (32,)
batch 1: x (32, 2)  y (32,)
batch 2: x (32, 2)  y (32,)
batch 3: x (32, 2)  y (32,)
batch 4: x (32, 2)  y (32,)

Three lines — pick, fetch, glue — are the whole loop. Everything DataLoader adds is industrialisation of those three lines.

The walkthrough

default_collate is the assembler. Given a list of (tensor, tensor) tuples, it stacks the first elements into one tensor and the seconds into another. Four samples of shape (2,) became one batch of (4, 2): samples stack along a new first dimension, which is why models expect batch-first input.

shuffle=True swaps the default in-order sampler for a random-permutation sampler. Every epoch — every full pass — deals a fresh permutation, so no batch is ever the same twice.

generator=... pins the shuffle to a seed you control, making the batch sequence reproducible without touching global seeds. For debugging a specific misbehaving batch, this is gold.

drop_last answers the awkward-remainder question: 160 samples at batch 32 divides evenly, but at batch 50 the last batch would hold 10. drop_last=True discards it — worth doing when a stray small batch breaks something like BatchNorm at batch size 1.

The loop restarting is free: each for over the loader builds a fresh iterator, asks the sampler for a new permutation, and goes again. An epoch is nothing more than one exhaustion of that iterator.

What batch_size silently changes

Batch size is not only speed. The gradient each step averages over the batch, so bigger batches mean smoother, fewer updates per epoch, and often a different best learning rate. When you change batch_size, expect to retune the learning rate — the classic pairing is scaling both together.

Common mistakes

Shuffling the validation loader. Harmless-looking, but it makes per-sample error inspection irreproducible between runs. Convention: shuffle=True for training, False everywhere else.

Iterating the dataset instead of the loader. for x, y in ds: works (Python falls back on __getitem__) but delivers unshuffled singles — models "train" suspiciously slowly and badly. Check what your loop iterates.

Assuming batch order equals dataset order. After shuffling, batch 0 is a random handful. Anything that depends on sequence — time series windows — must be handled inside the Dataset, not assumed from the loader.

Rebuilding the loader every epoch. Works, but wasteful once workers exist (lesson four). Build once, iterate many times.

Try it yourself

Set batch_size=50 and run with drop_last=False, then True. Count batches and check the final batch's shape in each case. Then print the first batch's y twice in two separate runs — same generator seed, same batch — to see reproducible shuffling.

What to learn next

Researcher — Mathematics and papers.

The composition, named precisely

DataLoader(ds, batch_size=32, shuffle=True) composes: RandomSampler(ds) → BatchSampler(sampler, 32, drop_last) → _MapDatasetFetcher → default_collate. Each is public and injectable: sampler= replaces index generation (weighted sampling), batch_sampler= replaces grouping (length-bucketing lives here), collate_fn= replaces assembly, and batch_size=None disables auto-batching entirely for datasets that emit ready-made batches. shuffle=True is exactly sugar for sampler=RandomSampler(ds) — there is no other magic in it.

RandomSampler draws a fresh permutation per iterator via torch.randperm on the loader's generator; sampling with replacement is an option, changing epoch semantics from "each sample once" to i.i.d. draws — a distinction that matters for convergence analyses, since most SGD theory assumes i.i.d. sampling while practice overwhelmingly uses without-replacement passes (see Bottou, 2009, Curiously Fast Convergence of some Stochastic Gradient Descent Algorithms, for why shuffling tends to beat i.i.d. draws empirically).

default_collate's recursion

The collator is type-driven and recursive: tensors → torch.stack (one new leading dim); numbers → 0-d tensors, then stacked; NumPy arrays → converted, then treated as tensors; mappings → collated per key, returning a dict of batched values; sequences → transposed and collated per position. So a dataset returning {"image": t, "meta": {"id": i}} yields a batch with the same nesting, every leaf batched. The recursion's failure mode is shape disagreement at any leaf — the stack expects each tensor to be equal size error dissected in the collate lesson.

A performance subtlety: when workers are active, collate runs in the worker, and the batch tensor is written into shared memory for zero-copy transfer to the parent; with pin_memory=True a separate pinning thread then stages it for async GPU copies (the starvation lesson). Collate cost is therefore parallelised, while anything done after the loader yields is not.

Epoch boundaries and iterator lifecycle

Each iter(loader) constructs a fresh iterator object: single-process (_SingleProcessDataLoaderIter) or multiprocess, which spawns workers, index queues, and a result queue. In-order delivery is guaranteed — the parent reorders worker results by sequence number, so nondeterministic worker scheduling never reorders batches; determinism of content thus reduces to sampler seeding plus per-worker RNG seeding. Worker lifecycle costs at this boundary motivate persistent_workers, measured with real numbers in the workers lesson.

Reading

  • The torch.utils.data documentation's "Data Loading Order and Sampler" and "Multi-process data loading" sections — normative.
  • Bottou (2009) — shuffled passes versus i.i.d. sampling.
  • Paszke et al. (2019), PyTorch: An Imperative Style... — the loader's design constraints.

What to learn next