Build a Custom PyTorch Dataset and DataLoader for Your Own Image Folder

Bundled datasets make PyTorch look easy in a way that does not survive contact with your own work. One line downloads sixty thousand perfectly labelled images, already the right size, already the right format, and you get to skip an entire category of work without noticing it exists. Then a real project arrives. The images are in a folder someone shared with you, some are the wrong size, a few have a transparency layer, and the labels live in the folder names or a spreadsheet. That gap is where most people learning PyTorch quietly stall. This tutorial closes it by building the class that sits between whatever data you actually have and the tidy batches a model expects.

What a Dataset Class Is For

PyTorch needs to ask two questions about your data: how many samples are there, and what is sample number 500? Anything that can answer both can be trained on. The Dataset class is simply the shape those two answers take.

That is a deliberately small contract, and the smallness is the point. Because PyTorch asks so little, almost anything can satisfy it. Images in folders, rows in a CSV, audio clips, records in a database, lines in a log file. You write a short class per data source, and everything downstream stays the same.

Why Bundled Datasets Hide This Work

The first three parts of this series used MNIST, which arrives through a single function call already labelled, already uniformly sized, already converted to tensors. That convenience quietly skips an entire category of engineering, and the skipped work is exactly what stalls people on their first real project.

Where Dataset Ends and DataLoader Begins

These two are constantly confused, and the split is worth stating plainly because it determines which one you need to modify when something is wrong.

A Dataset deals with one sample at a time. It knows where your files are, how to open one, and what label belongs to it. It has no concept of batches, shuffling, or parallelism.

The DataLoader deals with delivery. It decides which samples to fetch, groups them into batches, stacks them into a single tensor, and can load them in parallel background processes. It knows nothing about what your data means.

Why Transforms Live Inside the Sample Lookup

One design detail explains a lot of later behaviour. Transforms do not run when the dataset is created. They run every time a sample is fetched.

That sounds like wasted effort, and for fixed operations like resizing it partly is. But it is precisely what makes random augmentation possible. If a rotation is applied when a sample is read, then reading the same file in two different epochs produces two different images, and the model effectively trains on a larger dataset than you actually own.

It also explains the most common quiet mistake in this area, which is applying random augmentation to a validation set. Doing so makes your score change between runs on identical weights, so you lose the ability to tell whether a change helped.

Step 1: Generate a Dataset You Control

Rather than downloading anything, create 450 labelled images of coloured shapes. Reproducible sample data means you can verify your pipeline before pointing it at data that matters.

import os, random
from PIL import Image, ImageDraw

random.seed(7)
ROOT, CLASSES = "shapes", ["circle", "square", "triangle"]

for split, n in {"train": 120, "val": 30}.items():
    for cls in CLASSES:
        os.makedirs(f"{ROOT}/{split}/{cls}", exist_ok=True)
        for i in range(n):
            img = Image.new("RGB", (128, 128), (random.randint(230, 255),) * 3)
            d = ImageDraw.Draw(img)
            colour = (random.randint(0, 180), random.randint(0, 180), random.randint(0, 180))
            m = random.randint(15, 30)
            box = [m, m, 128 - m, 128 - m]
            if cls == "circle":
                d.ellipse(box, fill=colour)
            elif cls == "square":
                d.rectangle(box, fill=colour)
            else:
                d.polygon([(64, m), (128 - m, 128 - m), (m, 128 - m)], fill=colour)
            img.save(f"{ROOT}/{split}/{cls}/{cls}_{i:03d}.png")

print("Generated", sum(len(f) for _, _, f in os.walk(ROOT)), "images")

Generated 450 images
  train/circle: 120 images
  train/square: 120 images
  train/triangle: 120 images
  val/circle: 30 images
  val/square: 30 images
  val/triangle: 30 images

The resulting layout is the standard convention, where the folder name is the label:

shapes/
|-- train/
|   |-- circle/     120 images
|   |-- square/     120 images
|   `-- triangle/   120 images
`-- val/
    |-- circle/      30 images
    |-- square/      30 images
    `-- triangle/    30 images

Step 2: Write the Dataset Class

Three methods. That is the whole interface.

import torch
from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset, DataLoader

class ShapeDataset(Dataset):
    def __init__(self, root, transform=None):
        self.root = Path(root)
        self.transform = transform
        self.classes = sorted([d.name for d in self.root.iterdir() if d.is_dir()])
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
        self.samples = []
        for cls in self.classes:
            for path in sorted((self.root / cls).glob("*.png")):
                self.samples.append((path, self.class_to_idx[cls]))

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

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

raw = ShapeDataset("shapes/train")
print("Classes      :", raw.classes)
print("class_to_idx :", raw.class_to_idx)
print("Length       :", len(raw))
img, label = raw[0]
print("Sample type  :", type(img).__name__)
print("Sample label :", label)

Classes      : ['circle', 'square', 'triangle']
class_to_idx : {'circle': 0, 'square': 1, 'triangle': 2}
Length       : 360
Sample type  : Image
Sample label : 0

Four design decisions in that class are worth copying into your own work.

__init__ indexes paths, it does not load images. Building a list of file paths is cheap. Loading 360 images into memory is not, and for a dataset of any real size it is impossible. Files are opened only when a sample is actually requested.

sorted() appears twice, deliberately. Filesystem iteration order is not guaranteed to be stable across machines or runs. Sorting both the class list and the file list means circle maps to 0 on your laptop and on the training server, which matters enormously when you load a saved model and interpret its outputs.

.convert("RGB") normalises the channel count. PNG files may carry an alpha channel and some images are greyscale. Without this line you get a mix of three and four channel tensors, and batching fails partway through an epoch rather than immediately.

transform is stored, not applied. The transform runs inside __getitem__, which is what makes random augmentation work correctly. The next steps show why.

Step 3: Add Transforms

A transform pipeline is a list of operations applied in order to each sample.

from torchvision import transforms

tfm = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

ds = ShapeDataset("shapes/train", transform=tfm)
img, label = ds[0]
print("Tensor shape :", img.shape)
print("Tensor dtype :", img.dtype)
print("Value range  :", f"{img.min():.3f} to {img.max():.3f}")

Tensor shape : torch.Size([3, 64, 64])
Tensor dtype : torch.float32
Value range  : -1.467 to 2.379

Order matters and is not arbitrary.

The value range of negative 1.467 to 2.379 surprises people who expected 0 to 1. That is Normalize doing its job: it subtracts a mean and divides by a standard deviation, which centres values around zero and lets some fall below it. Those particular numbers are the channel statistics of ImageNet, and you use them whenever your model was pretrained on ImageNet, which is the subject of part five in this series.

Step 4: Wrap It in a DataLoader

The Dataset returns one sample. The DataLoader turns that into batches.

loader = DataLoader(ds, batch_size=32, shuffle=True, num_workers=0)

print("Batches per epoch:", len(loader))
batch_x, batch_y = next(iter(loader))
print("Batch images     :", batch_x.shape)
print("Batch labels     :", batch_y.shape)
print("Labels in batch  :", batch_y[:12].tolist())

Batches per epoch: 12
Batch images     : torch.Size([32, 3, 64, 64])
Batch labels     : torch.Size([32])
Labels in batch  : [1, 1, 0, 2, 2, 2, 1, 0, 0, 0, 0, 1]

360 samples in batches of 32 gives 12 batches, the last holding only 8 samples. Pass drop_last=True if a partial final batch causes problems, which it occasionally does with batch normalisation on very small remainders.

The shuffled labels confirm the loader is mixing classes rather than serving all circles, then all squares. That matters: an unshuffled loader would give the model 120 consecutive circles, and it would learn to predict circle for everything before ever seeing a square.

Verify the loader covers the full dataset exactly once per epoch:

counts = torch.zeros(len(ds.classes))
for _, ys in loader:
    counts += torch.bincount(ys, minlength=len(ds.classes))
for c, n in zip(ds.classes, counts.tolist()):
    print(f"  {c}: {int(n)}")

  circle: 120
  square: 120
  triangle: 120

Every sample appears once. Shuffling changes order, never membership.

Step 5: Separate Training and Validation Pipelines

This is the decision that most often causes trouble later, and it is easy to get right once you see it stated plainly.

train_tfm = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
])

val_tfm = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
])

train_ds = ShapeDataset("shapes/train", transform=train_tfm)
val_ds   = ShapeDataset("shapes/val",   transform=val_tfm)
print("train samples:", len(train_ds), "| val samples:", len(val_ds))

train samples: 360 | val samples: 90

Now confirm that the random transforms really do produce a different tensor each time the same index is read:

torch.manual_seed(1); first  = train_ds[0][0]
torch.manual_seed(2); second = train_ds[0][0]
print("Augmentation on, same index, identical tensors? ", torch.equal(first, second))
print("Augmentation off, same index, identical tensors?", torch.equal(val_ds[0][0], val_ds[0][0]))

Augmentation on, same index, identical tensors?  False
Augmentation off, same index, identical tensors? True

Index 0 points at the same file both times, yet the training pipeline returns different tensors because the flip and rotation are re rolled on each access. That is the mechanism behind augmentation: the model effectively sees a larger, more varied dataset without you storing a single extra file.

The validation pipeline is deterministic, and it has to be. If your validation transform randomised, your accuracy score would change between runs on identical weights, and you would have no way to tell whether a change you made helped or whether you simply got a luckier set of rotations.

Step 6: Know When ImageFolder Is Enough

When your data already sits in class named folders, torchvision does this for you.

from torchvision import datasets

ds = datasets.ImageFolder("shapes/train", transform=tfm)
print(ds.classes, len(ds))

Writing the class once is worth it precisely because the first row describes tutorials and the other four describe real projects.

Break It On Purpose

RuntimeError: stack expects each tensor to be equal size,
but got [3, 32, 32] at entry 0 and [3, 33, 32] at entry 1

Error Decoder

Performance Notes Worth Knowing Early

num_workers controls how many subprocesses load data in parallel. The default of 0 loads on the main process, which means the GPU sits idle while images are read and decoded. Setting it to 4 is a reasonable starting point on most machines. Note that on Windows and in Jupyter notebooks, values above 0 sometimes require guarding your entry point with if __name__ == "__main__":.

pin_memory=True speeds up the transfer from CPU to GPU and is worth setting whenever you train on a GPU. It has no benefit on CPU only training.

Keep __getitem__ light. It runs once per sample per epoch, so a dataset of 100,000 images calls it 100,000 times every epoch. Expensive parsing, network requests, or database queries inside that method will dominate your training time. Where possible, do heavy preprocessing once and cache the result to disk.

Conclusion

The Dataset class is small on purpose. Index your samples in __init__, report the count in __len__, load and return one sample in __getitem__, and the rest of PyTorch works with whatever data you have.

Four things are worth carrying forward. Index paths rather than loading files up front, or memory becomes the constraint. Sort your class list so label indices stay identical across machines. Put every random augmentation on the training pipeline and nothing random on validation. Resize inside the pipeline so every sample stacks cleanly into a batch.

Once this class exists, the model code from part one runs unchanged against your own images. That is the whole point of the abstraction, and it is why writing it once pays back across every project that follows.

FAQS