Datasets & DataLoaders

By the end of this lesson you will be able to wrap any data, from a NumPy array to a folder of images, in a PyTorch Dataset, and serve it to your model in shuffled mini-batches with a DataLoader. You will also load a real image dataset with torchvision and transform it on the fly.

In Training Loops, Batches & Epochs you wrote your own shuffle-and-slice loop. PyTorch splits that job into two small pieces, and almost every project uses them.

Two jobs, two classes

Think of a college library in Jaipur.

• The catalogue knows how many books there are and can fetch book number 417 when asked. That is a Dataset: it answers len(dataset) and dataset[i].
• The librarian with a trolley picks books in a random order, loads 32 at a time and wheels them to your table. That is a DataLoader: it shuffles, groups into batches, and can use several helpers in parallel.

Keeping them separate means the same dataset can be served in different ways: shuffled for training, in order for testing, with different batch sizes.

Step 1: data already in memory with TensorDataset

When your features and labels already fit in memory as tensors, TensorDataset wraps them in one line. random_split divides a dataset into training and validation parts.

import torch
from torch.utils.data import TensorDataset, DataLoader, random_split
from sklearn.datasets import load_digits

X, y = load_digits(return_X_y=True)
X = torch.tensor(X / 16.0, dtype=torch.float32)
y = torch.tensor(y)                                    # int64 labels

dataset = TensorDataset(X, y)
print("examples:", len(dataset))
features, label = dataset[0]
print("one example:", features.shape, "label", label.item())

generator = torch.Generator().manual_seed(42)           # repeatable split
train_set, val_set = random_split(dataset, [0.8, 0.2], generator=generator)
print("train:", len(train_set), " validation:", len(val_set))

train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
val_loader = DataLoader(val_set, batch_size=256, shuffle=False)
xb, yb = next(iter(train_loader))
print("one batch:", xb.shape, yb.shape, " batches per epoch:", len(train_loader))

examples: 1797
one example: torch.Size([64]) label 0
train: 1438  validation: 359
one batch: torch.Size([32, 64]) torch.Size([32])  batches per epoch: 45

dataset[0] returns a pair: one example's features and its label. The loader turns those pairs into batches: features stacked into (32, 64) and labels into (32,). len(train_loader) is the number of batches per epoch. Validation uses shuffle=False and a bigger batch, since no learning happens there.

Step 2: a custom Dataset

For anything else (CSV rows, image files, text) you write a small class with three methods:

1. __init__: load or locate the data, once.
2. __len__: return how many examples there are.
3. __getitem__(i): return example i as tensors, applying any transform.

Here is a dataset for a used-phone price table from a shop in Kochi, as it might come from a CSV file. It scales the features inside __getitem__.

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

class PhoneDataset(Dataset):
    def __init__(self, rows):
        self.rows = rows                                  # list of (ram_gb, age_years, price_k)

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

    def __getitem__(self, i):
        ram, age, price = self.rows[i]
        x = torch.tensor([ram / 16, age / 5], dtype=torch.float32)   # simple scaling
        y = torch.tensor([price], dtype=torch.float32)
        return x, y

rows = [(8, 1, 18.5), (4, 3, 7.0), (6, 2, 11.0), (12, 1, 27.0), (3, 4, 4.5)]
ds = PhoneDataset(rows)
print("length:", len(ds), " item 3:", ds[3])

loader = DataLoader(ds, batch_size=2, shuffle=False)
for xb, yb in loader:
    print([[round(v, 3) for v in row] for row in xb.tolist()], "->", yb.ravel().tolist())

length: 5  item 3: (tensor([0.7500, 0.2000]), tensor([27.]))
[[0.5, 0.2], [0.25, 0.6]] -> [18.5, 7.0]
[[0.375, 0.4], [0.75, 0.2]] -> [11.0, 27.0]
[[0.188, 0.8]] -> [4.5]

The last batch has only one example, since 5 does not divide evenly by 2. If a partial batch causes trouble (for example with batch normalisation), pass drop_last=True to the loader.

For images on disk, __init__ stores the list of file paths and __getitem__ opens one file. Loading files lazily, one at a time, means you never need the whole dataset in memory.

Step 3: torchvision datasets and transforms

The torchvision library ships ready-made datasets and image transforms: functions applied to each example as it is loaded. Here we load Fashion-MNIST, 70,000 small grey photos of clothes in 10 classes, which you will use in the project. It downloads about 30 MB the first time.

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import v2

transform = v2.Compose([
    v2.ToImage(),                                  # PIL image -> tensor image (C, H, W)
    v2.ToDtype(torch.float32, scale=True),         # 0..255 -> 0.0..1.0
    v2.Normalize(mean=(0.5,), std=(0.5,)),         # -> roughly -1..1
])
train_data = datasets.FashionMNIST("data", train=True, download=True, transform=transform)
test_data = datasets.FashionMNIST("data", train=False, download=True, transform=transform)
print("train:", len(train_data), " test:", len(test_data))
print("classes:", train_data.classes)

image, label = train_data[0]
print("one image:", image.shape, image.dtype, f"min {image.min():.1f} max {image.max():.1f}",
      "label:", train_data.classes[label])

loader = DataLoader(train_data, batch_size=64, shuffle=True)
images, labels = next(iter(loader))
print("batch:", images.shape, labels.shape)

train: 60000  test: 10000
classes: ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
one image: torch.Size([1, 28, 28]) torch.float32 min -1.0 max 1.0 label: Ankle boot
batch: torch.Size([64, 1, 28, 28]) torch.Size([64])

Each image arrives as (1, 28, 28): one grey channel, 28 × 28 pixels, already scaled. A batch adds the first dimension, giving (64, 1, 28, 28), the shape convolution layers expect. Training sets often get extra random transforms for data augmentation, such as v2.RandomHorizontalFlip(), while test sets get only the fixed conversion steps.

Step 4: the same idea in NumPy

A DataLoader is just a smart iterator. This small NumPy version shows the behaviour you rely on: a length, shuffling every epoch, and a smaller last batch unless you drop it.

import numpy as np

class SimpleLoader:
    def __init__(self, X, y, batch_size, shuffle, drop_last=False, seed=0):
        self.X, self.y, self.batch_size = X, y, batch_size
        self.shuffle, self.drop_last = shuffle, drop_last
        self.rng = np.random.default_rng(seed)

    def __len__(self):
        full, extra = divmod(len(self.X), self.batch_size)
        return full if self.drop_last or extra == 0 else full + 1

    def __iter__(self):
        order = self.rng.permutation(len(self.X)) if self.shuffle else np.arange(len(self.X))
        for start in range(0, len(order), self.batch_size):
            idx = order[start:start + self.batch_size]
            if self.drop_last and len(idx) < self.batch_size:
                return
            yield self.X[idx], self.y[idx]

X = np.arange(22).reshape(11, 2)
y = np.arange(11)
loader = SimpleLoader(X, y, batch_size=4, shuffle=True)
print("batches per epoch:", len(loader))
for epoch in range(2):
    print(f"epoch {epoch}:", [yb.tolist() for _, yb in loader])
print("drop_last:", len(SimpleLoader(X, y, 4, shuffle=False, drop_last=True)), "batches")

batches per epoch: 3
epoch 0: [[4, 6, 7, 2], [0, 3, 5, 10], [9, 8, 1]]
epoch 1: [[2, 10, 3, 6], [0, 8, 4, 9], [7, 5, 1]]
drop_last: 2 batches

Every epoch sees all 11 labels in a fresh order, in batches of 4, 4 and 3. With drop_last=True the incomplete batch is skipped.

Faster loading: workers and pinned memory

With images on disk, loading can become slower than training. Two DataLoader options help:

• num_workers=4 starts 4 background processes that prepare batches while the model trains. Start with the number of CPU cores, or fewer.
• pin_memory=True speeds up copying batches to an NVIDIA GPU.

On Windows and macOS, worker processes start by re-running your script, so the training code must sit under if __name__ == "__main__": when num_workers > 0. Without it you get an error about starting new processes before the current one has finished.

Common mistakes

• Shuffling the validation or test loader. It does no harm to the scores, but makes results harder to compare and debug. Use shuffle=False.
• Returning NumPy arrays or Python numbers from __getitem__. Return tensors with the right dtypes: float32 features and int64 class labels.
• Doing heavy work in __getitem__ that could be done once. Parse the CSV in __init__; do only per-example work in __getitem__.
• Random augmentation on the test set. Test transforms must be fixed.
• num_workers > 0 on Windows without the __main__ guard. Wrap the training code.

Key points

• A Dataset knows its length and returns one example by index; a DataLoader shuffles, batches and can load in parallel.
• TensorDataset wraps in-memory tensors; random_split makes training and validation sets.
• Custom datasets implement __init__, __len__ and __getitem__, loading files lazily.
• torchvision provides datasets such as Fashion-MNIST and transforms.v2 for conversion and augmentation.
• Use shuffle=True for training only, and add num_workers and pin_memory when loading is the bottleneck.

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

class SquaresDataset(Dataset):
    def __init__(self, n=20):
        self.values = torch.arange(1, n + 1, dtype=torch.float32)

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

    def __getitem__(self, i):
        x = self.values[i].reshape(1)
        return x, x ** 2

loader = DataLoader(SquaresDataset(), batch_size=6, shuffle=False, drop_last=True)
print("batches:", len(loader))
xb, yb = next(iter(loader))
print("x:", xb.ravel().tolist())
print("y:", yb.ravel().tolist())

batches: 3
x: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
y: [1.0, 4.0, 9.0, 16.0, 25.0, 36.0]