Custom Datasets for Real Data

The two-method contract

For any data shape you can reduce to "indexable list of samples", a custom Dataset just needs:

__len__(self) — total samples
__getitem__(self, idx) — return the sample at index idx

That's the entire contract. The DataLoader does the rest.

The patterns you'll write five hundred times

ImageFolder-style — a directory tree where subfolder names are class labels. torchvision ships ImageFolder for the canonical case; you'll write variants for "your data is similar but slightly different".
CSV / pandas — load a manifest file with paths and labels, fetch from disk on demand.
HDF5 / Parquet / WebDataset — for large datasets, read from a disk-backed format.
In-memory tensors — for small data that fits, just hold the tensors as attributes.

Lazy vs eager loading

Big datasets must load lazily — read from disk inside __getitem__, not in __init__. Small datasets can load eagerly and pre-process in __init__. The boundary is roughly "fits in RAM with headroom" vs "doesn't" — usually 10-100 GB depending on your machine.

Code

ImageFolder-style — load from a class-named directory tree·python

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

class ImageDirDataset(Dataset):
    """root/
        class_a/img1.jpg
                img2.jpg
        class_b/img3.jpg
    """
    def __init__(self, root, transform=None):
        self.root = root
        self.transform = transform
        self.samples = []
        self.classes = sorted(d for d in os.listdir(root)
                              if os.path.isdir(os.path.join(root, d)))
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}

        for cls in self.classes:
            cls_dir = os.path.join(root, cls)
            for fname in os.listdir(cls_dir):
                if fname.lower().endswith(('.jpg', '.jpeg', '.png')):
                    self.samples.append((os.path.join(cls_dir, fname), self.class_to_idx[cls]))

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

    def __getitem__(self, idx):
        path, label = self.samples[idx]
        img = Image.open(path).convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img, label

CSV-based dataset — manifest of paths + labels·python

import torch
import pandas as pd
from torch.utils.data import Dataset
from PIL import Image

class CSVImageDataset(Dataset):
    def __init__(self, csv_path, image_root, transform=None):
        self.df = pd.read_csv(csv_path)              # columns: 'path', 'label'
        self.image_root = image_root
        self.transform = transform

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img = Image.open(f"{self.image_root}/{row['path']}").convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img, int(row['label'])

Tabular Dataset — a pandas DataFrame of features·python

import torch
import pandas as pd
from torch.utils.data import Dataset

class TabularDataset(Dataset):
    def __init__(self, csv_path, target_col, feature_cols=None):
        df = pd.read_csv(csv_path)
        self.target = torch.tensor(df[target_col].values, dtype=torch.long)
        cols = feature_cols or [c for c in df.columns if c != target_col]
        self.features = torch.tensor(df[cols].values, dtype=torch.float32)

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

    def __getitem__(self, idx):
        return self.features[idx], self.target[idx]

Exercise

Write an ImageDirDataset for any small image set you have (or make 6 directories with 5 random JPEGs each). Confirm len(dataset), then sample dataset[0] and verify the image opens and the label matches the directory name. Time iterating through the whole dataset with num_workers=0 vs num_workers=4 — measure the speedup.