PyTorch · Lesson 5

Datasets & DataLoaders — Feeding the Loop

Real data doesn't arrive as one neat tensor. A Dataset says how to fetch one sample; a DataLoader batches and shuffles them for you. Together they're the standard intake pipe that feeds every training loop.

🎯 Mission: read & write real PyTorch
The one idea

PyTorch splits data handling into two clean jobs. A Dataset answers "how big are you?" (__len__) and "give me sample i" (__getitem__). A DataLoader wraps a dataset and turns it into a batched, shuffled iterable you loop over. Your training loop gains one inner line — for xb, yb in loader: — and otherwise stays exactly the same.

01 Why batches at all?

In Lessons 3–4 we shoved all 100 points through at once. That works for 100; it does not work for a million images that don't fit in memory. The fix is mini-batches: process a handful of samples, take a step, repeat.

Fits in memory

Only one batch is in memory/on the GPU at a time, so dataset size stops being bounded by RAM.

Learns faster

You take many small steps per pass instead of one big one — usually faster, smoother convergence (this is "mini-batch SGD").

Shuffling helps

Re-shuffling every epoch stops the model from learning the order of the data instead of the pattern.

💡 New vocabulary An epoch is one full pass over all batches. One step / iteration is one batch (one optimizer.step()). Batch size is how many samples per step — a key knob you'll see at the top of every script.

02 Dataset — how to get one sample

The simplest dataset wraps tensors you already have. TensorDataset pairs them up so that indexing returns one (features, target) sample.

from torch.utils.data import TensorDataset, DataLoader, random_split

ds = TensorDataset(X, y)          # X,y from the earlier lessons
print(len(ds))                  # 100        ← __len__
xi, yi = ds[0]                  # ONE sample ← __getitem__
print(xi.shape, yi.shape)        # torch.Size([1]) torch.Size([1])
💡 A dataset is "indexable + sized" Anything that supports len(ds) and ds[i] can be a dataset. Notice indexing returns a single sample (no batch dimension yet) — adding the batch dimension is the DataLoader's job.

Source: PyTorch — Datasets & DataLoaders tutorial.

03 DataLoader — batch + shuffle, for free

Wrap the dataset in a DataLoader, choose a batch_size, ask it to shuffle, and iterate. It stacks individual samples into a batch tensor with a new leading dimension.

loader = DataLoader(ds, batch_size=16, shuffle=True)
print(len(loader))               # 7   ← number of BATCHES (100 / 16, last partial)
xb, yb = next(iter(loader))
print(xb.shape, yb.shape)        # torch.Size([16, 1]) torch.Size([16, 1])
Dataset · 100 samples ds[i] → (xᵢ, yᵢ) __len__ · __getitem__ DataLoader shuffle, then cut into batches 7 batches of ≤16 training loop for xb, yb in loader: one step per batch batch tensor gains a leading dimension: (16, 1) — 16 samples stacked
Dataset = "one sample at a time." DataLoader = "batched, shuffled stream." The loop just iterates it.

04 The loop, now over batches

The only change to Lesson 4's loop is an inner for over the loader. Everything inside is the same five moves — they now run once per batch:

for epoch in range(31):
    for xb, yb in train_loader:    # ← the new line: iterate batches
        opt.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        opt.step()
epoch  0 | last-batch loss 4.67308
epoch 10 | last-batch loss 0.01068
epoch 30 | last-batch loss 0.00747
learned w=2.002, b=-3.000

Same line recovered, but now via many small batch-steps per epoch. Swap train_loader for one that streams a million images and not a single other line changes.

Run it: ~/projects/learn/public/courses/pytorch/practice/.venv/bin/python ~/projects/learn/public/courses/pytorch/practice/lesson05_dataloaders.py

05 Holding out a validation set

You must measure the model on data it didn't train on, or you can't tell learning from memorizing. random_split carves a dataset into pieces:

train_ds, val_ds = random_split(ds, [80, 20])   # 80 train, 20 validation
train_loader = DataLoader(train_ds, batch_size=16, shuffle=True)
val_loader   = DataLoader(val_ds,   batch_size=16)   # no shuffle needed for eval
⚠ Shuffle train, not val Shuffle the training loader (helps learning). Leave the validation/test loader unshuffled — order doesn't matter there, and not shuffling keeps evaluation reproducible. We'll actually use the validation set to measure accuracy in Lesson 6.

06 A custom Dataset — the real-data pattern

For real data (images on disk, rows of a CSV, tokenized text) you write your own Dataset. It's three methods — and it's the exact shape you'll read in countless repos:

from torch.utils.data import Dataset

class SquaresDataset(Dataset):
    def __init__(self, n):                 # set up / point at the data source
        self.x = torch.arange(n, dtype=torch.float32)
    def __len__(self):                     # how many samples
        return len(self.x)
    def __getitem__(self, i):              # fetch + prepare sample i → (input, label)
        return self.x[i], self.x[i] ** 2

sq = SquaresDataset(5)
print(len(sq), sq[3])     # 5 (tensor(3.), tensor(9.))
💡 This is where real data work lives __getitem__ is where you'd load an image from a path, apply transforms, tokenize a sentence, etc. — returning one ready-to-use sample. Hand that dataset to a DataLoader and batching, shuffling, and (optionally) parallel loading come for free.

Source: PyTorch — Datasets & DataLoaders.

07 Check yourself

What does indexing a Dataset (ds[i]) return?
A dataset has 100 samples and batch_size=16. What is len(loader)?
What's the leading dimension of a batch tensor from DataLoader(..., batch_size=16)?
Which three methods does a custom Dataset implement?
Which loader should you shuffle?

08 Flashcards

Q: What's the division of labour between Dataset and DataLoader?
A: Dataset = how to get one sample (__len__, __getitem__). DataLoader = wraps a dataset into a batched, shuffled, iterable stream.
Q: Why train on mini-batches instead of all data at once?
A: Data may not fit in memory; many small steps per epoch usually converge faster and smoother (mini-batch SGD); and per-batch shuffling avoids learning data order.
Q: epoch vs step (iteration)?
A: An epoch is one full pass over all batches. A step/iteration is one batch = one optimizer.step().
Q: What does a DataLoader add to a sample's shape?
A: A leading batch dimension. A sample of shape (1,) becomes a batch of shape (batch_size, 1).
Q: How do you make a train/validation split?
A: random_split(ds, [n_train, n_val]), then a DataLoader for each. Shuffle train, not val.
Q: Where does real data-loading logic go?
A: In a custom Dataset.__getitem__ — load the image/row/text for index i, apply transforms, and return (input, label).
👩‍🏫 I'm your teacher — ask me anything

Good asks: "what does num_workers do?", "how would a Dataset load images from a folder?", or "why is the last batch smaller and does it matter?" Next we put it all together: a real multi-layer classifier trained on your GPU. Say "go" for Lesson 6.