Custom Image Datasets and DataLoaders

Build the mental model

A Dataset class's job is simply to expose a contract for 'how do I get one sample': __len__() reports the total number of samples, and __getitem__(idx) defines how to return the (image, label) pair for a given index. Even when images are loaded from disk files (say, via PIL.Image.open), implementing this same interface is all that's required — it fully decouples where the data actually comes from from how the DataLoader consumes it. Wrapping synthetic in-memory tensors here skips file I/O entirely, but the interface is identical either way.

The DataLoader then consumes this Dataset and handles the practical concerns: batching, shuffling, and (with more worker processes) parallel loading. Setting shuffle=True reshuffles the sample order every epoch, which prevents the model from memorizing any sequential ordering in the data — for example, if the dataset happened to be sorted by class, training without shuffling could produce batches containing only one class each, destabilizing training. Because the DataLoader exposes a simple iterator interface (for batch in loader), the training loop itself can be written without caring how large the underlying dataset is.

Connect it to a real scenario

Suppose the Tutorial Platform is training an internal model to detect near-duplicate lesson screenshot previews. If the screenshot images have already been converted to in-memory tensors during an earlier preprocessing pipeline (rather than being re-read from disk each time), a custom Dataset class like this one can wrap those in-memory tensors directly, letting the DataLoader work with zero extra file I/O overhead.

A Dataset subclass only needs __len__ and __getitem__, decoupling the data source from the DataLoader.
DataLoader handles the practical concerns — batching, shuffling, and parallel loading.
shuffle=True reorders samples every epoch, preventing the model from memorizing any ordering in the data.

Try the working example

python

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

class FakeImageDataset(Dataset):
    def __init__(self, num_samples=100, num_classes=10):
        # Synthetic in-memory data: no files touched
        self.images = torch.randn(num_samples, 3, 32, 32)
        self.labels = torch.randint(0, num_classes, (num_samples,))

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

    def __getitem__(self, idx):
        return self.images[idx], self.labels[idx]

dataset = FakeImageDataset(num_samples=100, num_classes=10)
loader = DataLoader(dataset, batch_size=16, shuffle=True)

images, labels = next(iter(loader))
print(f"batch images shape: {tuple(images.shape)}")
print(f"batch labels shape: {tuple(labels.shape)}")
print(f"dataset length: {len(dataset)}")

You should see

Prints exactly: 'batch images shape: (16, 3, 32, 32)', 'batch labels shape: (16,)', 'dataset length: 100'. Since batch_size=16 was specified, the first batch contains 16 samples, and the dataset's total length is 100.

5-minute try-it

Add an optional transform parameter to FakeImageDataset and apply it inside __getitem__ before returning the image, if one was provided. Instantiate the dataset with a simple transform like lambda x: x * 2 and verify the returned image values change accordingly.

One important caution

Using self.images[idx:idx+1] instead of self.images[idx] inside __getitem__ keeps an extra batch dimension, causing DataLoader's default collate function to produce an unexpected extra dimension in the batch.

Setting batch_size larger than num_samples together with drop_last=True results in zero batches being produced, so next(iter(loader)) raises a StopIteration error.

PyTorch Tutorials — Writing Custom Datasets, DataLoaders and Transforms — Computer Vision