PyTorch Dataset

The Dataset class is PyTorch's standard interface for representing a data source — implement two methods, and PyTorch's data loading machinery (covered next) handles the rest.

The Minimal Interface

import torch
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

    def __len__(self):
        return len(self.features)           # total number of examples

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]   # ONE example, given its index

features = torch.randn(1000, 10)
labels = torch.randint(0, 2, (1000,))
dataset = MyDataset(features, labels)

print(len(dataset))         # 1000
print(dataset[0])            # (feature_vector, label) for the first example

Why Just These Two Methods Are Enough

PyTorch's DataLoader (next note) only needs to know how many examples exist (__len__) and how to fetch any single one by index (__getitem__) — everything else (shuffling, batching, parallel loading) is handled generically by DataLoader itself, regardless of what kind of data the specific Dataset subclass wraps.

Built-in Datasets

import torchvision.datasets as datasets
import torchvision.transforms as transforms

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
mnist = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
print(len(mnist))         # 60,000
image, label = mnist[0]

Common Mistakes

• Loading and preprocessing the entire dataset eagerly inside __init__ when it doesn't fit comfortably in memory — for large datasets, __getitem__ should typically load each item lazily (e.g. reading an image file from disk on demand), not all upfront.
• Returning inconsistent types or shapes from __getitem__ across different indices — every call must return data in a compatible format so DataLoader can correctly batch multiple items together.

Interview Relevance

Q: "What two methods must a custom PyTorch Dataset implement, and why are these two sufficient?" __len__ (returning the total number of examples) and __getitem__ (returning a single example given its index). These two methods are sufficient because PyTorch's DataLoader handles all the generic machinery — shuffling, batching, parallel loading — on top of this simple interface, regardless of the underlying data source or format.

Practice Question

Why is it often better for __getitem__ to load an image from disk on demand, rather than loading every image into memory inside __init__, for a dataset with millions of images?

Related DL Notes

• PyTorch DataLoader
• PyTorch Custom Datasets
• Mini-Batch Gradient Descent