01 · Dataset Basics 📇

4 min read

Everything in PyTorch's data pipeline starts with one small, elegant contract: "tell me how many samples you have, and give me the one at index i."

1. The 'Why'

In Phase 1 your data was a handful of tensors you typed out by hand. Real datasets don't work that way. You might have 60,000 images in a folder, a million rows in a CSV, or audio clips scattered across subdirectories with labels stored somewhere else entirely. You can't just load all of it into one giant tensor — it might not fit in memory, and even when it does, you need to fetch, transform, and shuffle individual samples on demand. What you need is a uniform way to represent a dataset so the rest of PyTorch can consume it without caring where the data actually lives.

That uniform representation is the Dataset. It's a deliberately tiny interface: a dataset only has to answer two questions — how many samples are there? and what is sample number i? By implementing just those two methods, your custom data source instantly becomes compatible with the entire PyTorch ecosystem, most importantly the DataLoader (sub-module 02) that handles batching and shuffling for you. This separation is powerful: your Dataset focuses purely on reading and returning one sample, and PyTorch handles everything else. Nail this contract and every data source you'll ever meet — images, text, tabular, audio — fits the same clean mold.

2. Core Concepts

The Dataset contract. Subclass torch.utils.data.Dataset and implement two methods: __len__ (returns the number of samples) and __getitem__(idx) (returns the sample at that index, typically a (features, label) tuple).

Map-style vs. iterable-style. The common kind is map-style: it supports indexing (dataset[i]), which is what __len__/__getitem__ provide. There's also an IterableDataset for streaming data that can't be indexed (e.g. a live feed); you'll rarely need it early on.

Return tensors, not raw objects. __getitem__ should return tensors (or things easily collated into tensors). Do any file reading and conversion inside __getitem__ so each sample is model-ready when it comes out.

Lazy loading. A good Dataset loads each sample only when asked inside __getitem__, rather than loading everything up front. This keeps memory use low no matter how big the dataset is.

3. Code in Action

A minimal custom Dataset

import torch
from torch.utils.data import Dataset

class ToyDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features        # e.g. a tensor of shape (N, D)
        self.labels = labels            # e.g. a tensor of shape (N,)

    def __len__(self):
        return len(self.features)       # how many samples exist

    def __getitem__(self, idx):
        # return the single sample at position idx as (X, y)
        return self.features[idx], self.labels[idx]

X = torch.randn(100, 4)                 # 100 samples, 4 features each
y = torch.randint(0, 3, (100,))         # 100 integer labels in {0,1,2}
ds = ToyDataset(X, y)

print(len(ds))                          # 100 -> calls __len__
sample_x, sample_y = ds[0]              # calls __getitem__(0)
print(sample_x.shape, sample_y)         # torch.Size([4]) tensor(...)

A lazy, file-backed Dataset (the realistic pattern)

from torch.utils.data import Dataset
from PIL import Image
import torch

class ImageFolderDataset(Dataset):
    def __init__(self, file_paths, labels, transform=None):
        self.file_paths = file_paths    # list of paths on disk
        self.labels = labels            # matching list of labels
        self.transform = transform      # optional preprocessing (sub-module 03)

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

    def __getitem__(self, idx):
        img = Image.open(self.file_paths[idx]).convert("RGB")  # load ONLY this file
        if self.transform:
            img = self.transform(img)   # e.g. resize + to-tensor + normalize
        label = torch.tensor(self.labels[idx])
        return img, label

Indexing behaves like a normal sequence

print(ds[5])          # the 6th sample
print(ds[-1])         # the last sample (negative indexing works)
xs = [ds[i][0] for i in range(3)]   # grab the first three feature tensors

4. Common Pitfalls

Pitfall 1 — Loading everything in __init__. Reading all files into memory up front defeats the purpose and crashes on large datasets. Store lightweight references (paths, indices) in __init__ and do the actual loading lazily inside __getitem__.

Pitfall 2 — Returning inconsistent shapes or types. The DataLoader will try to stack samples into a batch, which fails if __getitem__ sometimes returns a float32 tensor and sometimes float64, or images of different sizes. Make every sample come out the same shape and dtype (transforms help — sub-module 03).

Pitfall 3 — Off-by-one in __len__. If __len__ reports more samples than exist, __getitem__ will eventually get an out-of-range index and crash mid-epoch. Make sure __len__ matches your data exactly.

5. Further Reading & Watch List

➡️ Next: 02 · The DataLoader