DL S016 · Datasets, DataLoaders, and Collate

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 16 / 80 · Module M03 · ~2 hours

The story — the day Andrej saw his GPU at 4%

Andrej Karpathy gave a talk at Stanford in 2019 where he shared the single most-forwarded slide of his Tesla years: a screenshot of nvidia-smi showing eight V100 GPUs at 4% utilization during a Tesla Autopilot training run. Every one of those GPUs cost $10K. They were running at the speed of a Raspberry Pi. The bottleneck wasn't the model. It wasn't the network. It was that the data pipeline was decoding raw camera frames on a single Python process, and the GPU spent 96% of every batch cycle sitting there waiting. Two engineers spent a week rewriting the loader with proper multiprocessing, prefetching, and pinned memory. Training time dropped from 3 weeks to 3 days. The model architecture never changed.

That's the lesson: your beautiful new architecture is running on a $10,000 GPU that can crunch a batch in 3 milliseconds — and it's sitting idle 90% of the time because your data pipeline is decoding a JPEG on a single CPU thread. On a toy problem you'd never notice. On a real training run it turns a 3-hour job into a 30-hour job.

PyTorch's answer is the Dataset / DataLoader split. Dataset says "here is item i". DataLoader says "here are batch_size random items, prefetched using num_workers processes, moved to pinned memory, ready when the GPU is ready". This decoupling is what lets a research student and Meta both write essentially the same training loop.

Today we do all of it: MNIST out of the box, then a custom Dataset for your own data, then the collate_fn trick that lets you handle variable-length text (this is what unlocks NLP), then the war stories around workers and pinned memory that will save you the "why is my GPU only 20% utilized" week.

Prerequisites

• Session 014 — tensors.
• Session 015 — you have a model waiting to be fed.

1 · The two-piece design

Every PyTorch data pipeline is exactly two objects.

Dataset — knows nothing about batching. Just:

• __len__(self) → total number of items
• __getitem__(self, i) → the i-th item, usually as a (x, y) tuple

DataLoader — wraps a Dataset. Knows about batching, shuffling, parallelism.

from torch.utils.data import DataLoader
 
loader = DataLoader(dataset,
                    batch_size=128,
                    shuffle=True,
                    num_workers=4,
                    pin_memory=True)
 
for xb, yb in loader:
    # xb.shape == (128, ...)
    ...

That's it. Every model in this series will consume data this way.

2 · MNIST in three lines

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
 
train_ds = datasets.MNIST("./data", train=True, download=True,
                          transform=transforms.ToTensor())
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=2)
 
for xb, yb in train_loader:
    print(xb.shape, yb.shape)     # torch.Size([128, 1, 28, 28]) torch.Size([128])
    break

datasets.MNIST downloads on first use, caches to ./data, returns a Dataset whose items are (PIL_Image, int_label). transforms.ToTensor() converts PIL → (C, H, W) float tensor in [0, 1]. DataLoader handles the rest.

Same pattern for CIFAR-10, CIFAR-100, FashionMNIST, ImageNet, etc.

3 · Writing a custom Dataset

The moment you have your own data (a folder of your own images, a CSV of tabular data, a JSONL of text), you write a Dataset.

from torch.utils.data import Dataset
from PIL import Image
import os, torch
 
class ImageFolderDataset(Dataset):
    def __init__(self, root, transform=None):
        self.samples = []                        # list of (path, label)
        self.transform = transform
        for label, cls in enumerate(sorted(os.listdir(root))):
            cls_dir = os.path.join(root, cls)
            for fn in os.listdir(cls_dir):
                self.samples.append((os.path.join(cls_dir, fn), label))
 
    def __len__(self):
        return len(self.samples)
 
    def __getitem__(self, i):
        path, label = self.samples[i]
        img = Image.open(path).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, label

A few design rules worth staring at.

4 · The DataLoader parameter tour

Every one of these will bite you at some point. Read them all now, once, so you recognize the symptoms later.

4.1 batch_size

Number of items per batch. Trade-off: bigger = smoother gradients + better GPU utilization; smaller = more updates per epoch + fits smaller GPUs. Doubling batch size ~roughly doubles memory. If it OOMs, halve it. Session 024 shows gradient accumulation for when 1 is too many.

4.2 shuffle=True/False

Randomize order each epoch. Almost always True for training, always False for validation (so metrics are comparable across epochs and reproducible). Uses a fresh random permutation each epoch.

4.3 num_workers=N

Number of subprocess workers to prefetch batches. 0 = load in the main process (blocks GPU). 2–8 is typical. Rule of thumb: start with num_workers = min(8, os.cpu_count()). Higher isn't always better — too many workers thrash disk and RAM.

4.4 pin_memory=True

Copies each batch into pinned (page-locked) RAM. This lets .to("cuda", non_blocking=True) do an async DMA transfer — the CPU issues the copy and moves on to loading the next batch while the previous one is still moving to GPU. Always True when training on GPU. No-op on CPU.

4.5 drop_last=True

Drops the last incomplete batch. Set True if your model architecture assumes a fixed batch size (rare) or if a partial batch would trigger batch norm weirdness with 1 sample. Otherwise False.

4.6 persistent_workers=True

Keeps worker processes alive across epochs instead of respawning. Big win if __init__ of the dataset is even slightly expensive. Default False for backward compat; you almost always want True.

4.7 prefetch_factor=2 (default)

Each worker prefetches this many batches ahead. Increase if disk is slow and workers keep waiting.

5 · collate_fn — the escape hatch that unlocks NLP

DataLoader calls collate_fn(list_of_items) to turn N items into one batch. The default (default_collate) does the sensible thing: it stacks tensors along a new dim-0. Works perfectly for images (all same shape). Fails immediately for text, where every sentence has a different length.

5.1 The problem

sentences = [
    torch.tensor([1, 2, 3]),           # len 3
    torch.tensor([4, 5, 6, 7, 8]),     # len 5
    torch.tensor([9]),                 # len 1
]
torch.stack(sentences)
# RuntimeError: stack expects each tensor to be equal size

5.2 The fix — pad to max length in batch

from torch.nn.utils.rnn import pad_sequence
 
def collate_text(batch):
    """batch is a list of (token_ids, label)."""
    tokens = [b[0] for b in batch]
    labels = torch.tensor([b[1] for b in batch])
    # pad_sequence pads with 0 on the right to max length in batch
    padded = pad_sequence(tokens, batch_first=True, padding_value=0)
    return padded, labels
 
loader = DataLoader(text_dataset, batch_size=32, collate_fn=collate_text)

pad_sequence returns (batch, max_len) with 0s where sentences were shorter. Every transformer training loop you'll ever look at has some version of this pattern. When we get to LLMs in M07, this line is what makes variable-length batching possible.

5.3 Bonus — dynamic batching by length

Advanced but worth mentioning: bucket sentences by length before batching, so you don't pad a 5-token sentence to the length of a 500-token one in the same batch. That's what torch.utils.data.BatchSampler + a custom sampler is for. We revisit this in M07.

6 · Debugging the "GPU is only 20% utilized" problem

Classic symptom, classic diagnosis.

Step 1: Time one batch, model-only. Put a dummy xb, yb on the GPU and loop through the model without a DataLoader. If this is fast, your model isn't the bottleneck.

Step 2: Time one epoch with num_workers=0. If it's 5× slower than num_workers=4, data is the bottleneck.

Step 3: Profile the dataset. Wrap __getitem__ and time it. Common culprits:

• Decoding JPEGs from disk on every call (fix: use fast image formats like WebP, or cache resized versions)
• Heavy augmentation on CPU (fix: move augmentations to GPU with kornia, or reduce num_ops)
• Small __getitem__ items with high per-call overhead (fix: batch items server-side, or use a WebDataset)

Step 4: Verify pin_memory=True and non_blocking=True in .to().

xb = xb.to(device, non_blocking=True)

Both together let CPU→GPU transfer overlap with computation. Free perf.

7 · War stories

def __getitem__(self, i):
    if not hasattr(self, "_conn"):
        self._conn = sqlite3.connect(self.path)
    return self._conn.execute("SELECT ... WHERE id=?", (i,)).fetchone()

8 · Diagram — the pipeline

Every arrow between "Dataset" and "Model" is an opportunity for a bottleneck. num_workers widens the top; pin_memory + non_blocking widens the bottom.

8.5 · IterableDataset and streaming — when your data doesn't fit

Everything so far assumed a map-style dataset: you can ask for item i and get it. That works for MNIST (60K images) and even ImageNet (1.3M). It stops working for the 15-trillion-token corpus used to train Llama 3, or the multi-terabyte video dataset used for Sora-style models. You can't hold the index of every item in RAM, and you can't seek to item i in a compressed tar stream.

Enter torch.utils.data.IterableDataset. Instead of __getitem__, you implement __iter__ — the dataset is a stream. This is what powers modern large-scale training.

from torch.utils.data import IterableDataset
import tarfile, io, torch
 
class ShardedTarDataset(IterableDataset):
    def __init__(self, shard_paths):
        self.shard_paths = shard_paths
 
    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        # partition shards across workers so each sees a disjoint set
        shards = self.shard_paths[info.id::info.num_workers] if info else self.shard_paths
        for path in shards:
            with tarfile.open(path) as tar:
                for member in tar:
                    yield decode(tar.extractfile(member).read())

Two new concerns appear:

1. Worker sharding — with map-style datasets, the DataLoader partitions indices across workers automatically. With iterable, you must partition, or every worker will yield the same items. Use torch.utils.data.get_worker_info().
2. No shuffle=True — DataLoader can't shuffle a stream. Do it inside __iter__ with a shuffle buffer (read N ahead, yield randomly from the buffer).

The de facto standard for large-scale training is WebDataset (pip install webdataset) — an IterableDataset that reads sharded .tar files, does shuffle buffers, handles worker sharding, and streams from cloud storage (S3, GCS) transparently. Meta uses it. Stability AI uses it. If your dataset is >100 GB, use WebDataset.

8.6 · StatefulDataLoader — resumable training (PyTorch 2.4+, torchdata 0.8)

One of the ugliest bugs in long-running training: your job crashes 3 days in. You restart from a checkpoint. Great — but your DataLoader starts over from batch 0. If shuffle is random, you re-see items you already trained on this epoch. If shuffle isn't random, you re-see the same first 1000 batches again and skew the loss.

StatefulDataLoader (from torchdata, integrated into PyTorch's data ecosystem as of 2024) fixes this. Its state_dict() captures where in the epoch you were; load_state_dict() resumes from exactly that point. This is now the standard for pretraining runs on Llama-scale models.

from torchdata.stateful_dataloader import StatefulDataLoader
 
loader = StatefulDataLoader(dataset, batch_size=128, num_workers=4)
# during training, alongside model+optimizer state:
checkpoint['loader'] = loader.state_dict()
# on resume:
loader.load_state_dict(checkpoint['loader'])

A note on TorchData's twisty history: the older DataPipe API (2022–2023) was deprecated in mid-2024 in favor of a simpler design centered on StatefulDataLoader. If you find IterDataPipe in tutorials, know it's legacy — use plain IterableDataset + StatefulDataLoader today. See the PyTorch 2.4 release notes for the migration.

8.7 · HuggingFace datasets — the pragmatic answer

For 90% of NLP/vision work in 2025, you won't write a Dataset from scratch. You'll do:

from datasets import load_dataset
ds = load_dataset("HuggingFaceFW/fineweb-edu", split="train", streaming=True)
for row in ds:
    ...

HuggingFace's datasets library gives you Arrow-backed memory-mapped access, streaming from the Hub, on-the-fly filtering/mapping, and a .with_format("torch") that turns rows into tensors. It's what almost every LLM paper's data code actually looks like. We'll use it heavily starting in M07.

The rule of thumb for choosing:

• Your own image folder, <10 GB → custom Dataset (§3).
• Your own text/tabular data → datasets.Dataset.from_dict or .from_pandas.
• Public dataset → load_dataset("...") from the Hub.

• 100 GB streaming → WebDataset or load_dataset(..., streaming=True).

9 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §6 (the four-step diagnosis) and §5.2 (pad_sequence) in 24 hours. Full session on day 7.

Next session (S017): the training loop you'd trust in production — seeds, checkpoint/resume (now including StatefulDataLoader!), TensorBoard/W&B logging, gradient monitoring, torch.compile gotchas, and the "log everything for future you" mindset.

Sticky note (keep on your desk): pin_memory=True + .to(dev, non_blocking=True) = free GPU perf. pad_sequence(..., batch_first=True) in collate_fn unlocks NLP. If GPU util <50%, blame data first, model second.

Previous: ← DL S015 · nn.Module · Next: DL S017 · The Training Loop →