02 · The DataLoader 🚚

5 min read

A Dataset gives you one sample at a time. The DataLoader turns that trickle into the batched, shuffled, parallel stream a training loop actually needs.

1. The 'Why'

Your Dataset from sub-module 01 knows how to produce a single sample, but training a neural network almost never happens one sample at a time. We train on mini-batches — small groups of samples processed together — because batching makes far better use of the parallel hardware in a GPU and because averaging the gradient over a batch gives a more stable learning signal than a single noisy example. We also need to shuffle the data every epoch so the model doesn't learn the accidental order of the file system, and we'd like to load the next batch while the GPU is busy on the current one, so the expensive accelerator never sits idle waiting on disk. Doing all of this by hand — slicing indices, randomizing, spawning worker processes — would be tedious and bug-prone.

The DataLoader does all of it for you. You hand it a Dataset and a few settings, and it becomes an iterable that yields ready-to-use batches: it groups samples, stacks them into batched tensors, reshuffles each epoch, and can spin up background worker processes to prefetch data in parallel. In practice, the DataLoader is the object your training loop iterates over directly. Understanding its handful of key arguments — batch_size, shuffle, num_workers — is what makes your training both correct and fast.

2. Core Concepts

batch_size. How many samples per batch. The DataLoader fetches that many items from the Dataset and stacks them, adding a leading batch dimension (e.g. (32, 4) for 32 samples of 4 features).

shuffle. When True, the order is randomized every epoch. Use True for training and False for validation/test (where order doesn't matter and reproducibility is nice).

num_workers. How many background processes load data in parallel. 0 means loading happens in the main process (simple, but can bottleneck the GPU). A few workers often speeds things up dramatically; the best value depends on your machine.

collate_fn. The function that merges a list of samples into a batch. The default stacks tensors, which works when every sample has the same shape. For variable-length data (like sentences), you supply a custom collate_fn to pad them.

drop_last. If the dataset size isn't divisible by batch_size, the final batch is smaller. Set drop_last=True to discard it when a uniform batch size matters.

3. Code in Action

Wrapping a Dataset in a DataLoader

import torch
from torch.utils.data import Dataset, DataLoader

class ToyDataset(Dataset):
    def __init__(self, n=100):
        self.X = torch.randn(n, 4)
        self.y = torch.randint(0, 3, (n,))
    def __len__(self):  return len(self.X)
    def __getitem__(self, i):  return self.X[i], self.y[i]

ds = ToyDataset()
loader = DataLoader(
    ds,
    batch_size=16,     # 16 samples per batch
    shuffle=True,      # reshuffle every epoch (use for training)
    num_workers=2,     # 2 background processes prefetch data
    drop_last=False,   # keep the final (possibly smaller) batch
)

Iterating batches (the shape you'll see in training)

for batch_X, batch_y in loader:            # each loop yields one batch
    print(batch_X.shape, batch_y.shape)    # torch.Size([16, 4]) torch.Size([16])
    break                                  # (just peeking at the first batch)

# One full pass over `loader` = one EPOCH. The training loop (Module 05)
# wraps this in an outer "for epoch in range(...)" loop.
print("Batches per epoch:", len(loader))   # ceil(100 / 16) = 7

A custom `collate_fn` for variable-length data

import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader

def pad_collate(batch):
    # batch is a list of (sequence_tensor, label) with sequences of different lengths
    seqs, labels = zip(*batch)
    padded = pad_sequence(seqs, batch_first=True)  # pad to the longest in the batch
    return padded, torch.stack(labels)

# loader = DataLoader(text_dataset, batch_size=8, collate_fn=pad_collate)

4. Common Pitfalls

Pitfall 1 — Shuffling the validation/test set (or not shuffling training). Shuffle training data so the model doesn't learn ordering artifacts; keep validation/test unshuffled for reproducible, comparable metrics. Swapping these is a subtle but common mistake.

Pitfall 2 — Cranking num_workers too high (especially on Windows/notebooks). More workers isn't always faster and can cause memory blow-ups or, on Windows, errors if the DataLoader isn't guarded under if __name__ == "__main__":. Start with 0 or 2 and tune up.

Pitfall 3 — Forgetting to move each batch to the device. The DataLoader yields CPU tensors. Inside the loop you must send both the batch and the model to the same device: batch_X = batch_X.to(device). Skipping this triggers the "tensors on different devices" error from Module 01.

5. Further Reading & Watch List

⬅️ Prev: 01 · Dataset Basics · ➡️ Next: 03 · Transforms