PyTorch Custom Datasets and DataLoaders

Real-World Scenario

An ML engineer trains an image classifier on 500,000 product images stored on disk. Loading them all into RAM crashes the machine. A custom Dataset reads each image only when it’s requested and applies augmentation on the fly, while a DataLoader with several worker processes prepares the next batches in parallel so the GPU isn’t left waiting.

The Dataset Protocol

A map-style dataset needs two methods: __len__ and __getitem__. The DataLoader handles everything else.

import torch
from torch.utils.data import Dataset, DataLoader

class Squares(Dataset):
    def __init__(self, n):
        self.n = n

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        x = torch.tensor([float(idx)])
        return x, x ** 2                     # one (input, target) pair

ds = Squares(10)
print(len(ds), ds[3])

loader = DataLoader(ds, batch_size=4)
for xb, yb in loader:
    print(xb.shape, yb.flatten().tolist())

$ python squares.py
10 (tensor([3.]), tensor([9.]))
torch.Size([4, 1]) [0.0, 1.0, 4.0, 9.0]
torch.Size([4, 1]) [16.0, 25.0, 36.0, 49.0]
torch.Size([2, 1]) [64.0, 81.0]

ds[3] calls __getitem__(3) and returns one (input, target) pair. The DataLoader asks for indices in order, then its default collate function stacks the individual (1,) tensors into a (4, 1) batch. Ten items with batch_size=4 leave a final batch of two.

Loading Images Lazily from Disk

First, a small dataset on disk — six solid grey images and a CSV of labels:

import csv
from pathlib import Path
from PIL import Image

root = Path("shapes")
(root / "images").mkdir(parents=True, exist_ok=True)
rows = []
for i in range(6):
    label = "bright" if i % 2 else "dark"
    shade = 220 if label == "bright" else 30
    Image.new("RGB", (64, 48), (shade, shade, shade)).save(root / "images" / f"img{i}.png")
    rows.append({"image_path": f"img{i}.png", "label": label})

with open(root / "labels.csv", "w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["image_path", "label"])
    writer.writeheader()
    writer.writerows(rows)
print(open(root / "labels.csv").read().strip())

$ python make_images.py
image_path,label
img0.png,dark
img1.png,bright
img2.png,dark
img3.png,bright
img4.png,dark
img5.png,bright

The dataset keeps only the CSV in memory and opens each image inside __getitem__, so memory use doesn’t grow with the number of images:

from pathlib import Path
import pandas as pd
import torchvision.transforms as T
from PIL import Image
from torch.utils.data import Dataset, DataLoader

class ImageCSVDataset(Dataset):
    def __init__(self, csv_path, image_dir, transform=None):
        self.df = pd.read_csv(csv_path)              # only paths and labels in memory
        self.image_dir = Path(image_dir)
        self.transform = transform
        self.classes = sorted(self.df["label"].unique())
        self.label2idx = {name: i for i, name in enumerate(self.classes)}

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        image = Image.open(self.image_dir / row["image_path"]).convert("RGB")  # read on demand
        if self.transform:
            image = self.transform(image)
        return image, self.label2idx[row["label"]]

transform = T.Compose([T.Resize((32, 32)), T.RandomHorizontalFlip(), T.ToTensor()])
ds = ImageCSVDataset("shapes/labels.csv", "shapes/images", transform=transform)
print(ds.classes, ds.label2idx)

images, labels = next(iter(DataLoader(ds, batch_size=4, shuffle=False)))
print(images.shape, images.dtype, labels.tolist())
print([round(v, 3) for v in images.mean(dim=(1, 2, 3)).tolist()])

$ python image_dataset.py
['bright', 'dark'] {'bright': 0, 'dark': 1}
torch.Size([4, 3, 32, 32]) torch.float32 [1, 0, 1, 0]
[0.118, 0.863, 0.118, 0.863]

Class ids come from the sorted label names, so bright is 0 and dark is 1 — alphabetical, not the order they appear in the file. ToTensor converts the PIL image to a float32 tensor shaped (channels, height, width) with values scaled to [0, 1]: the dark images (pixel value 30) average 30/255 = 0.118 and the bright ones 220/255 = 0.863. Random augmentations like RandomHorizontalFlip run again every time an item is loaded, so each epoch sees slightly different versions of the same images.

Splitting and Batching

import torch
from torch.utils.data import TensorDataset, DataLoader, random_split

X = torch.randn(10_000, 20)
y = torch.randint(0, 5, (10_000,))
dataset = TensorDataset(X, y)

train_ds, val_ds = random_split(dataset, [0.8, 0.2],
                                generator=torch.Generator().manual_seed(42))
print(len(train_ds), len(val_ds), type(train_ds).__name__)

for drop_last in (False, True):
    loader = DataLoader(train_ds, batch_size=256, shuffle=True, drop_last=drop_last)
    sizes = [len(yb) for _, yb in loader]
    print(f"drop_last={drop_last}: {len(loader)} batches, last has {sizes[-1]}")

$ python split.py
8000 2000 Subset
drop_last=False: 32 batches, last has 64
drop_last=True: 31 batches, last has 256

random_split accepts fractions and returns Subset objects — views that hold indices into the original dataset, not copies of the data. Passing a seeded generator makes the split reproducible. 8,000 samples in batches of 256 leave a final batch of only 64; drop_last=True discards it, which avoids a tiny, noisy last step and keeps every batch the same size for layers like BatchNorm.

Weighted Sampling for Class Imbalance

import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler

rng = np.random.default_rng(42)
y = torch.tensor(rng.choice([0, 1], 10_000, p=[0.9, 0.1]))
dataset = TensorDataset(torch.arange(10_000), y)        # row index stands in for features

class_counts = torch.bincount(y)
sample_weights = (1.0 / class_counts)[y]                # rare class -> larger weight
sampler = WeightedRandomSampler(sample_weights, num_samples=len(y), replacement=True,
                                generator=torch.Generator().manual_seed(0))
loader = DataLoader(dataset, batch_size=64, sampler=sampler)

idx, labels = map(torch.cat, zip(*loader))
print("class counts:          ", class_counts.tolist())
print(f"original positive rate  {y.float().mean():.1%}")
print(f"sampled positive rate   {labels.float().mean():.1%}")
print("positive draws:        ", int(labels.sum()))
print("distinct positive rows:", len(idx[labels == 1].unique()))

$ python weighted.py
class counts:           [9015, 985]
original positive rate  9.8%
sampled positive rate   50.1%
positive draws:         5005
distinct positive rows: 979

Each sample’s weight is one over its class count, so the two classes carry equal total weight and a batch drawn from the sampler is about half positive. Because the sampler draws with replacement, the positive rate rises to 50% by repeating rows: 5,005 positive draws came from only 979 distinct positive samples, each seen about five times per epoch. That repetition is a real overfitting risk, and the model’s predicted probabilities no longer reflect the true 10% base rate. A sampler replaces shuffle; passing both raises an error.

Custom Collate for Variable-Length Data

The default collate function stacks tensors, which only works when they all have the same shape.

import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader

sentences = [torch.tensor([5, 12, 3]), torch.tensor([8, 2]), torch.tensor([1, 15, 9, 4])]
labels = [1, 0, 1]
data = list(zip(sentences, labels))

try:
    next(iter(DataLoader(data, batch_size=3)))
except RuntimeError as e:
    print("default collate:", e)

def pad_collate(batch):
    seqs, ys = zip(*batch)
    lengths = torch.tensor([len(s) for s in seqs])
    return pad_sequence(seqs, batch_first=True, padding_value=0), lengths, torch.tensor(ys)

padded, lengths, ys = next(iter(DataLoader(data, batch_size=3, collate_fn=pad_collate)))
print(padded)
print(lengths.tolist(), ys.tolist())

$ python collate.py
default collate: stack expects each tensor to be equal size, but got [3] at entry 0 and [2] at entry 1
tensor([[ 5, 12,  3,  0],
        [ 8,  2,  0,  0],
        [ 1, 15,  9,  4]])
[3, 2, 4] [1, 0, 1]

Sentences of lengths 3, 2, and 4 cannot be stacked, and the error says so. A collate_fn receives the list of samples for one batch and can build the batch any way it likes — here, padding to the longest sequence with pad_sequence and returning the real lengths as well, ready for pack_padded_sequence in an RNN.

Streaming with IterableDataset

When data arrives as a stream — a huge file read line by line, a database cursor — subclass IterableDataset and implement __iter__ instead of __getitem__. Worker processes need extra care:

import torch
from torch.utils.data import IterableDataset, DataLoader

class Numbers(IterableDataset):
    def __init__(self, n):
        self.n = n

    def __iter__(self):
        return iter(range(self.n))           # every worker runs this in full

if __name__ == "__main__":
    for workers in (0, 2):
        loader = DataLoader(Numbers(6), batch_size=None, num_workers=workers)
        print(f"num_workers={workers}:", sorted(int(x) for x in loader))

$ python stream_dup.py
num_workers=0: [0, 1, 2, 3, 4, 5]
num_workers=2: [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5]

With two workers, every number arrives twice. Each worker gets its own copy of the dataset and runs __iter__ from the start; nothing tells them to split the work. get_worker_info() fixes that by telling each worker who it is:

import torch
from torch.utils.data import IterableDataset, DataLoader, get_worker_info

class ShardedNumbers(IterableDataset):
    def __init__(self, n):
        self.n = n

    def __iter__(self):
        info = get_worker_info()
        if info is None:                     # main process: everything
            return iter(range(self.n))
        return iter(range(info.id, self.n, info.num_workers))   # every k-th item

if __name__ == "__main__":
    for workers in (0, 2, 3):
        loader = DataLoader(ShardedNumbers(6), batch_size=None, num_workers=workers)
        print(f"num_workers={workers}:", sorted(int(x) for x in loader))

$ python stream_shard.py
num_workers=0: [0, 1, 2, 3, 4, 5]
num_workers=2: [0, 1, 2, 3, 4, 5]
num_workers=3: [0, 1, 2, 3, 4, 5]

get_worker_info() returns None in the main process and an object with id and num_workers inside a worker. Taking every num_workers-th item from offset id gives each worker a disjoint shard, and the combined output has every number exactly once for any worker count. The if __name__ == "__main__": guard matters whenever num_workers > 0: on platforms that start workers by re-importing the script, it stops each worker from starting its own loaders.

When Workers Help

import time
import torch
from torch.utils.data import Dataset, DataLoader, TensorDataset

class SlowDisk(Dataset):
    """Pretends each item costs 2 ms to read and decode, like an image on disk."""
    def __len__(self):
        return 1024
    def __getitem__(self, idx):
        time.sleep(0.002)
        return torch.randn(3, 32, 32), idx % 10

def seconds(dataset, workers):
    t0 = time.perf_counter()
    for _ in DataLoader(dataset, batch_size=64, num_workers=workers):
        pass
    return time.perf_counter() - t0

if __name__ == "__main__":
    in_memory = TensorDataset(torch.randn(1024, 3, 32, 32), torch.randint(0, 10, (1024,)))
    for workers in (0, 2, 4):
        print(f"workers={workers}  slow disk {seconds(SlowDisk(), workers):.2f}s"
              f"   in memory {seconds(in_memory, workers):.2f}s")

$ python loader_speed.py
workers=0  slow disk 2.24s   in memory 0.01s
workers=2  slow disk 1.18s   in memory 0.09s
workers=4  slow disk 0.64s   in memory 0.10s

When each item is slow to produce — here a 2 ms pause stands in for reading and decoding an image — workers overlap that waiting, and the time drops roughly in proportion to their number. For data already in memory, the opposite happens: num_workers=0 is fastest, because workers add process start-up and copying every batch between processes for no benefit. Set num_workers for disk-bound datasets and measure; don’t assume more is faster.

Practice

import torch
from torch.utils.data import DataLoader

data = list(range(10))

def order(seed):
    loader = DataLoader(data, batch_size=5, shuffle=True,
                        generator=torch.Generator().manual_seed(seed))
    return [b.tolist() for b in loader]

print(order(0))
print(order(0))
print(order(1))

$ python p1_reproducible.py
[[3, 7, 5, 2, 0], [8, 1, 6, 9, 4]]
[[3, 7, 5, 2, 0], [8, 1, 6, 9, 4]]
[[4, 6, 9, 2, 3], [1, 0, 7, 8, 5]]

import torch
from torch.utils.data import Dataset

class Squares(Dataset):
    def __init__(self, n, transform=None):
        self.n, self.transform = n, transform

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        x = torch.tensor([float(idx)])
        y = x ** 2
        if self.transform:
            x = self.transform(x)          # transform the input, not the target
        return x, y

ds = Squares(10, transform=lambda x: x / 10)
print(ds[3], ds[9])

$ python p2_transform.py
(tensor([0.3000]), tensor([9.])) (tensor([0.9000]), tensor([81.]))

import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader

data = [(torch.tensor([5, 12, 3]), 1), (torch.tensor([8, 2]), 0), (torch.tensor([1, 15, 9, 4]), 1)]

def masked_collate(batch):
    seqs, ys = zip(*batch)
    padded = pad_sequence(seqs, batch_first=True, padding_value=-1)
    return padded, padded != -1, torch.tensor(ys)

padded, mask, ys = next(iter(DataLoader(data, batch_size=3, collate_fn=masked_collate)))
print(padded)
print(mask.int())
print(mask.sum(dim=1).tolist())

$ python p3_mask.py
tensor([[ 5, 12,  3, -1],
        [ 8,  2, -1, -1],
        [ 1, 15,  9,  4]])
tensor([[1, 1, 1, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 1]], dtype=torch.int32)
[3, 2, 4]

from torch.utils.data import IterableDataset, DataLoader

class Numbers(IterableDataset):
    def __iter__(self):
        return iter(range(6))

try:
    DataLoader(Numbers(), batch_size=2, shuffle=True)
except ValueError as e:
    print(e)

$ python p4_iterable_shuffle.py
DataLoader with IterableDataset: expected unspecified shuffle option, but got shuffle=True