Writing a custom Dataset

The Dataset contract

torch.utils.data.Dataset is a small contract. A map-style dataset implements two methods: __len__ returns how many samples exist, and __getitem__(idx) returns one sample, usually a (features, label) pair of tensors. Anything that follows this contract can be fed to a DataLoader.

Keep __init__ light. Store a list of file paths or row offsets, and do the expensive work (opening an image, decoding audio, tokenising text) inside __getitem__. That way a dataset larger than your RAM still works, and the loading work can be spread across worker processes.

Return tensors with the dtypes the model expects: float32 features and int64 class labels. For data already in memory, TensorDataset(X, y) saves you writing a class. For streams with no fixed length, such as reading from a socket, subclass IterableDataset and implement __iter__ instead. torch.utils.data.random_split(ds, [0.8, 0.2], generator=...) splits a dataset into training and validation subsets.

How a sample reaches the model

Files on disk are read one sample at a time by the Dataset's getitem, transforms turn each sample into a normalised tensor, the DataLoader groups samples into batches, and the model consumes the batches.

A Dataset for exam scores

Each row is three test scores and a pass/fail label. The scores are scaled to [0, 1] and returned as float32; the label is int64, ready for CrossEntropyLoss.

import torch
from torch.utils.data import Dataset

class StudentMarks(Dataset):
    """Predict pass (1) or fail (0) from three test scores."""

    def __init__(self, rows):
        self.rows = rows                    # list of (scores, label)

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

    def __getitem__(self, idx):
        scores, label = self.rows[idx]
        x = torch.tensor(scores, dtype=torch.float32) / 100.0
        y = torch.tensor(label, dtype=torch.long)
        return x, y

ds = StudentMarks([([72, 65, 80], 1), ([30, 41, 25], 0), ([55, 60, 48], 1)])
print(len(ds))      # 3
x, y = ds[1]
print(x, y)         # scores 0.30, 0.41, 0.25 and label 0

Reproducible splits

random_split shuffles with the global random generator unless you pass one. Use random_split(ds, [0.8, 0.2], generator=torch.Generator().manual_seed(42)) so the validation set is the same every run. Otherwise, comparing two experiments compares two different validation sets.

Quick check: Which two methods must a map-style `Dataset` implement?

• `__iter__` and `__next__`
• `load` and `transform`
• `forward` and `backward`
• `__len__` and `__getitem__`