Dataset Objects

The Dataset abstraction

PyTorch's torch.utils.data.Dataset is the canonical contract for "a thing you can index into to get one example." Subclass it, implement __len__ and __getitem__, and the rest of the data pipeline (loaders, samplers, augmentation) just works on top.

The Dataset is where you do per-example transformation: read a row from disk, decode an image, tokenize a string, normalize a tensor. Anything that depends on the example index lives here. Anything that depends on a batch (collation, padding to max length) lives in the DataLoader's collate function.

Tip: Datasets are lazy by default — __getitem__ only runs when something asks for an example. This is what lets you train on terabyte-scale data without loading any of it into RAM upfront.

Map-style vs iterable-style

Map-style (most common): subclass Dataset, implement __getitem__(idx). Random access, supports shuffling, length is known. Best for finite, indexable datasets.

Iterable-style: subclass IterableDataset, implement __iter__. Sequential streaming, no random access, length may be unknown. Best for log streams, distributed sharding, or huge web-scale corpora.

Three datasets you should be able to write blindfolded

Tensor wrapper — TensorDataset(X, y) for in-memory data. Fine for prototyping.
Image folder — torchvision.datasets.ImageFolder('path/') for class-per-subfolder image datasets.
Custom CSV / JSONL — read a manifest in __init__, lazy-load each example in __getitem__.

Principle: Most data bugs hide in __getitem__. Print one example before you start training. Verify the shape, dtype, and label match what your model expects.

Code

A custom Dataset for an image manifest·python

import torch
from torch.utils.data import Dataset
from PIL import Image
import csv, pathlib

class ManifestImageDataset(Dataset):
    def __init__(self, manifest_path, root, transform=None):
        self.root = pathlib.Path(root)
        self.transform = transform
        with open(manifest_path) as f:
            self.rows = list(csv.DictReader(f))  # [{'filename': ..., 'label': ...}]

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

    def __getitem__(self, idx):
        row = self.rows[idx]
        img = Image.open(self.root / row["filename"]).convert("RGB")
        if self.transform is not None:
            img = self.transform(img)
        label = int(row["label"])
        return img, label

The Dataset abstraction

Map-style vs iterable-style

Three datasets you should be able to write blindfolded

Code

External links

Exercise

Progress

Comments 0