Dataset and DataLoader
Lesson 5 of 8
5Dataset and DataLoader
The big pictureA model does not see all your data at once; it sees batches. A Dataset knows how to hand you one example, and a DataLoader groups examples into batches, shuffles them and loads them in parallel. Batch size is not a technical detail: it changes how the model learns.
What the picture shows: One epoch is cut into batches, and each batch is one training step. Five examples in batches of two give three batches, and the last holds one example unless you ask for drop_last=True.
Map of this part
The concepts in order, click any one to jump to it:
Two ways into the same idea
You do not feed the model all the data at once, just as you do not eat a whole plate in one bite. You split it into bites (batches), shuffle the order each time so the model does not memorise it, and prepare the next bite while it chews the current one. That is what Dataset and DataLoader do.
A Dataset defines __len__ and __getitem__ to return one example, and a DataLoader groups examples into batch_size batches with shuffle=True and num_workers for parallel loading. Bigger batches are faster with smoother gradients, but use more memory and may need a larger learning rate.
Before You Start
⚙️A training loop by hand
You know each training step works on a batch, and an epoch is one full pass over the training data.
🧊Tensors
You know the first dimension of a shape is usually the batch size, like (2, 1) here.
The Core Idea
A model does not see all the data at once; it sees batches. In this lesson you build a Dataset that hands out one example, then a DataLoader that groups them into batches, then see what batch size does, then why we shuffle.
1Dataset: one example at a time
A Dataset is any class with two methods: __len__ says how many examples you have, and __getitem__ returns example i as tensors. Look at the picture: a number goes in, a ready example comes out.
In __init__ you keep only the list of paths, so the whole dataset never sits in memory. Random image augmentation also happens inside __getitem__.
2DataLoader: batches, shuffling, parallel loading
A DataLoader turns a Dataset into batches, and shuffles the order every epoch when you ask for shuffle=True. Look at the picture: each epoch starts with a new order, then goes through the batches one after another.
This code needs PyTorch on your machine, so the output shown is recorded from a real run. Install it with pip install torch and run it yourself.
import torch
from torch.utils.data import DataLoader, Dataset
class Reviews(Dataset):
def __init__(self, texts, labels):
self.texts = texts
self.labels = labels
def __len__(self):
return len(self.texts)
def __getitem__(self, i):
length = torch.tensor([len(self.texts[i])], dtype=torch.float32)
return length, torch.tensor(self.labels[i])
data = Reviews(["ممتاز", "سيئ جدا", "رائع وسريع", "لا أنصح", "جيد"], [1, 0, 1, 0, 1])
print("size:", len(data))
print("item 1:", data[1])
torch.manual_seed(0)
for x, y in DataLoader(data, batch_size=2, shuffle=True):
print(tuple(x.shape), y.tolist())
Five examples in batches of two, so the last batch holds one. The feature is crude here (text length) because the point is the shape: each batch is (2, 1) features with a list of labels. In practice __getitem__ is where you read an image from disk or tokenise a sentence.
With num_workers=4, four processes prepare the next batches while the GPU trains on the current one. And drop_last=True drops a small final batch if it bothers you.
3Batch size changes the gradient
A small batch means more steps and a noisier gradient. A large batch means a more accurate gradient, fewer steps and more memory. Look at the picture: how much the same model's gradient varies between batches.
We compute the same model's gradient (at weight zero) on every batch in one epoch, and measure how much it varies between batches.
import numpy as np
rng = np.random.default_rng(0)
X = rng.uniform(-2, 2, size=1024)
y = 3 * X - 1 + rng.normal(scale=0.5, size=1024)
for batch_size in [8, 128, 1024]:
grads = []
for start in range(0, 1024, batch_size):
xb, yb = X[start:start + batch_size], y[start:start + batch_size]
grads.append(2 * ((0.0 * xb - yb) * xb).mean())
print(f"batch={batch_size:4} steps per epoch={len(grads):3} gradient spread={np.std(grads):.3f}")
With batch 8: 128 steps per epoch and every gradient differs a lot from the others (spread 2.8). With 128: eight steps and spread 0.69. With all the data: one perfectly accurate step. Noise in small batches is not always bad, since it sometimes helps the model escape flat regions, but it needs a smaller learning rate.
4Why we shuffle the training data
Sorted data, by class or by time, gives biased batches. Shuffling every epoch makes each batch a small picture of the whole dataset.
import random
labels = [0] * 6 + [1] * 6
batch_size = 4
print("no shuffle:")
for start in range(0, len(labels), batch_size):
print(" ", labels[start:start + batch_size])
random.seed(1)
order = list(range(len(labels)))
random.shuffle(order)
print("shuffled:")
for start in range(0, len(order), batch_size):
print(" ", [labels[i] for i in order[start:start + batch_size]])
Without shuffling, the first batch is all class 0 and the last all class 1. The model learns "always say 0", is then forced into "always say 1", and swings back and forth. After shuffling every batch holds both classes, so each gradient sees the whole picture.
Shuffle training only. Validation and test keep a fixed order, so results are comparable and every prediction maps back to its row.
How It Works, Step by Step: how data gets from disk to the model
The Dataset keeps only the list of paths in __init__, and reports their count from __len__.
__getitem__(i) reads example i from disk and turns it into tensors, and random augmentation happens here.
The DataLoader shuffles the order every epoch when shuffle=True, and groups examples into batches.
With num_workers above zero, separate processes prepare the next batches while the GPU trains on the current one.
Worked example
A PyTorch Dataset is any class with two methods: __len__ says how many examples you have, and __getitem__ returns example i as tensors. That is all a DataLoader needs.
__init__. It works with a thousand images and collapses with a hundred thousand. Keep only the list of paths in __init__ and read the image itself in __getitem__.
backward calls before one step, covered in the fine-tuning track.
Try it without code
Myth vs Reality
| Myth | Reality |
|---|---|
| Batch size is a technical detail set only by memory. | It changes how the model learns. With batch 8 the gradient spread was 2.8 and with 128 it was 0.69, and when you grow it a lot you usually need a larger learning rate. |
| Shuffling always helps, so shuffle everything. | Shuffle training only. Validation and test keep a fixed order so results are comparable and each prediction maps back to its row. |
| Slow training means the model is big. | Often the GPU is waiting for the disk. If its utilisation is low while the CPU is busy, the fix is num_workers and preparing images, not the model. |
Try it yourself
Which DataLoader should have shuffle=False?
Evaluation has no training steps affected by order, and a fixed order makes results comparable between runs and makes it easy to match each prediction to its original row.
GPU utilisation is only 20% during training and the CPU is at 100%. First thing to try?
The GPU is waiting for data. The CPU is the bottleneck, so spread loading over more processes, or pre-process images to their final size. The learning rate has nothing to do with speed.
Quick review cards
__len__ returning the number of examples and a __getitem__ returning example i as tensors. Heavy loading happens in the latter so not all the data sits in memory.num_workers, all without writing the loop yourself.Check Your Understanding
You doubled batch size from 32 to 64 and left everything else alone. Why might training take longer to reach the same loss even though each step is more accurate?
Because the number of steps per epoch halved, and a more accurate gradient does not make up for half the steps at the same learning rate. The well-known rule of thumb: when you double the batch, raise the learning rate by a similar factor. In practice very large batches also need a learning-rate warm-up so they do not blow up at the start.
Sorted data: six examples of class 0 then six of class 1, batches of 4, no shuffling. What happens to training?
The first batch is all 0 and the last all 1, so the model learns "always say 0", is then pushed to "always say 1", and swings between them. Shuffling puts both classes in every batch.
Resources for this lesson
Grouped by type. Start with the official docs for precision, or an article or video for the explanation.
Done with this lesson?
Sign in to save your progress.