# Writing a custom Dataset — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-dataset

> Implement a map-style Dataset with __len__ and __getitem__, load data lazily, and split it reproducibly.

## 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.

![Diagram of 5 connected stages from left to right: files on disk, Dataset __getitem__, transforms, DataLoader batches, model.](assets/figures/pytorch/section-6-map.svg) — Figure 6.1: files on disk -> Dataset __getitem__ -> transforms -> DataLoader batches -> model

## 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`.

```python
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.

**Quiz:** Which two methods must a map-style `Dataset` implement?

- [ ] `__iter__` and `__next__`
- [ ] `load` and `transform`
- [ ] `forward` and `backward`
- [x] `__len__` and `__getitem__`

*Answer:* `__len__` and `__getitem__`. Map-style datasets are indexed: `__len__` gives the size and `__getitem__` returns one sample. `__iter__` belongs to `IterableDataset`.