Data Management in PyTorch — Taming the Oxford 102 Flowers Dataset (v2)

01 — The Framing Problem

Same model, different results

Three teams built the flower classifier. Same layers, same parameters, wildly different accuracy.

It turned out, the model wasn't the variable. The data handling was.

A sophisticated model can't rescue a broken pipeline. If you can't access and prepare your data correctly, nothing downstream works.

What you're given

The Oxford 102 Flowers dataset zipped folder with JPEGs named image_00001.jpg, image_00002.jpg, and so on. The labels in a .mat file, MATLAB's binary format, mapping each image to one of 102 categories.

Let's see how we can handle this data effectively.

Three things that can go wrong

Access — finding the files and matching them to labels.
Quality — getting images into a size, type, and shape the model accepts.
Efficiency — loading in batches instead of one at a time.

Each has its own failure mode and we will tackle them in this post.

02 — Data Access

The custom Dataset class

Unlike common datasets like MNIST, this isn't pre-packaged. Let's see how we can teach Pytorch to read it

A Pytorch Dataset class asks three questions

Method	Answers
`__init__`	Where's the data?
`__len__`	How many samples?
`__getitem__`	Give me sample 42.

init

Setup reads the label file and stores paths. It doesn't open or load any images.

This is called lazy loading. Loading 8K+ images upfront costs gigabytes of RAM. The image is read only when someone asks for it.

The labels also run 1 to 102. PyTorch expects 0 to 101. Missing this will silently destablize training

len

Returns the number of samples

getitem

Given an index, returns the image and label. The image is read from disk. THe image file name is updated based on the new index

oxford_flowers_dataset.py

class OxfordFlowersDataset(Dataset):

    # Setup where to find images and labels
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.img_dir = os.path.join(root_dir, 'jpg')

        labels_mat = scipy.io.loadmat(os.path.join(root_dir, 'imagelabels.mat'))
        self.labels = labels_mat['labels'][0] - 1  # labels start at 1, pytorch expects 0
        self.transform = transform

    # Count number of total samples
    def __len__(self):
        return len(self.labels)

    # How to get image and label number 'idx'
    def __getitem__(self, idx):
        img_name = f"image_{idx+1:05d}.jpg"
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path)
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

Test before you build Pull dataset[0] the moment the class exists. The filename bug surfaces immediately here — and invisibly later.

test_it_early.py

dataset = OxfordFlowersDataset('flower_data')
print(f'Total Samples : {len(dataset)}')
img, label = dataset[0]
print(label)
img

output

Total Samples : 8189
76

8,189 samples, correctly linked. Label 76 for the first image, zero-indexed. The same three methods work for text, audio, anything — light setup, lazy loading.

03 — Quality Problems

Transform pipelines

The Dataset works. The DataLoader still crashes.

A batch is one tensor of shape (batch, channels, height, width). You can't stack images of different sizes into that. So the loader throws.

inspect_raw_samples.py

for i in [0, 100, 500]:
    img, labels = dataset[i]
    print(f'Image Size : {img.size}')
    print(f'Image Type : {type(img)}')

output

Image Size : (591, 500)   Image Type : PIL.JpegImagePlugin.JpegImageFile
Image Size : (588, 500)   Image Type : PIL.JpegImagePlugin.JpegImageFile
Image Size : (667, 500)   Image Type : PIL.JpegImagePlugin.JpegImageFile

Two problems. The sizes differ — different cameras, different years. And the type is wrong: PIL images in, tensors expected. Transforms fix both.

Solving Different size problem with transforms

Resize((224, 224)) — forces both dimensions and squashes anything rectangular.
Resize(256) — scales the shorter edge to 256, keeps aspect ratio.
CenterCrop(224) — takes a square from the center.

Original flower image, 591 by 500 pixels

original · (591, 500)

Resize(256) · (302, 256)

CenterCrop(224) · (224, 224)

Shorter edge to 256, aspect ratio intact. Then a clean 224 square.

standard_pipeline.py

transform = transforms.Compose([
    ## image transforms
    transforms.Resize(256),
    transforms.CenterCrop(224),

    ## tensor transforms
    transforms.ToTensor(),   
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
    std=[0.229, 0.224, 0.225])
])

Common Pitfall Don't use the name transforms; it will overwrite the module transforms itself

Solving PIL Image in dataset with `ToTensor()`

ToTensor() does the following things:

Converts PIL Image to PyTorch tensor
Moves channels from the back to the front
Divides each pixel by 255, scaling values to [0, 1]

Scaling to a common range makes 10% mean 10% everywhere. It also keeps the math stable — networks multiply constantly, otherwise errors in the 255s will compound fast.

Normalize spreads them out

Values between 0 and 1 can still bunch up. Bright images cluster near 1, dark ones near 0. Detail gets compressed into a narrow band and subtle differences disappear.

Normalize rescales by a mean and standard deviation so values spread evenly across the range. The numbers used here are standard ones, chosen to work well on natural images.

ToTensor is a one-way bridge

Before it, you hold an image. After it, a tensor. Some transforms only work on one side.

Modern TorchVision lets most image transforms work on either. Normalize does not — it needs a tensor and errors on an image. Order your Compose accordingly: image operations first, tensor operations after.

Debug one step at a time

When a pipeline misbehaves, apply the transforms individually and print the shape after each. You see exactly where the data stops looking right.

one_transform_at_a_time.py

img, _ = dataset[0]
print(img.size)                              # (591, 500)

resized = transforms.Resize(256)(img)
print(resized.size)                          # (302, 256)

cropped = transforms.CenterCrop(224)(resized)
print(cropped.size)                          # (224, 224)

tensored = transforms.ToTensor()(cropped)
tensored.shape                               # torch.Size([3, 224, 224])

output

torch.Size([4, 3, 224, 224])

The batch that crashed now stacks. Four images, three channels, 224 square, normalized.

04 — Efficiency

Splitting and batching

Why not train on everything? Because then you have no way to know if it works on anything new.

Split	Job
Training	What the model learns from. Seen every epoch.
Validation	Checks progress while you're still tuning.
Test	Used once, at the very end.

split.py

from torch.utils.data import random_split

train_size = int(0.7 * len(dataset))
val_size   = int(0.15 * len(dataset))
test_size  = len(dataset) - train_size - val_size   

train_dataset, val_dataset, test_dataset = random_split(
    dataset, [train_size, val_size, test_size]
)

output

Training   : 5732
Validation : 1228
Test       : 1229

random_split assigns images randomly. Without it, an ordered dataset puts all daisies in one split and all roses in another. Also nothing gets copied: you just get three views of the same underlying dataset.

Shuffle, and when not to

loaders.py

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader   = DataLoader(val_dataset,   batch_size=32, shuffle=False)
test_loader  = DataLoader(test_dataset,  batch_size=32, shuffle=False)

img, l = next(iter(train_loader))   # quick inspection without a loop
img.shape, l.shape

output

(torch.Size([32, 3, 224, 224]), torch.Size([32]))

Each iteration yields one batch: 32 images and their 32 labels. Use next(iter(loader)) to inspect one without writing a loop.

Shuffle the training set for two reasons. An ordered dataset lets the model learn position instead of features. And seeing only daisies for many batches makes it forget daisies once the roses start.

Validation and test don't shuffle. The model isn't learning from them, so neither problem applies. Shuffling changes only the serving order — the dataset itself is untouched.

Batches and epochs

5,732 training images at batch size 32 gives 179 full batches plus one batch of four. That last small batch is the remainder, and it's normal.

One epoch is all 180 batches, every image seen once. Ten epochs is each image seen ten times, in a different order each pass.

Two expensive mistakes

Reading a file inside __getitem__. A CSV opened per sample is reloaded once per sample — thousands of times per epoch, tens of thousands across a run. Read it once in __init__.
CUDA out of memory. Drop the batch size first. Start at 32 or 16, then climb.

05 — Bug-Proofing

Augmentation, errors, and a visual check

A pipeline that works and a pipeline that survives are different things.

Augmentation

Every training image was shot in similar conditions. A model trained only on centered flowers in good light fails on off-centered and badly lit images.

PyTorch allows some transforms that will allow you to show the same flower flipped, rotated, and re-lit, and the model will learn shape and color instead of position and lighting.

You could save the variants as files. PyTorch does better: it applies random transforms every time an image is loaded. Flower 42 is flipped in one epoch and darkened in the next. Unlimited variation, no extra storage.

RandomHorizontalFlip

RandomRotation(20)

ColorJitter(brightness)

Each transform alone, on the same crop.

two_transforms.py

train_transform = transforms.Compose([
    # Augmentation
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(degrees=20),
    transforms.ColorJitter(brightness=0.2),

    # Standard Preprocessing
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
    std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    # Standard Preprocessing only -- no random changes
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
    std=[0.229, 0.224, 0.225])
])

Validation gets no augmentation If the validation images change every epoch, you can't tell an improved model from a different input.

One bad file shouldn't end the run

Training crashes two hours in on a single corrupted image. Or on one that's valid but too small for the transforms. Four defenses:

Log, don't crash. An error_log list in __init__.
Check before using. verify() catches corruption, a size check catches tiny images, convert('RGB') catches grayscale.
Fall through. If it still fails, record what happened and return the next image.
Review afterward. One method to list what broke, so you can fix or drop it.

Is the augmentation too strong?

Push it too far and flowers become abstract blobs. A model that can't identify the subject can't learn its features.

The check is to look. Every call to dataset[idx] returns a new random variant, so render several side by side. Undo the normalization first, or you're looking at rescaled numbers instead of colors.

What you see	What it means
Recognizable every time	Good
All identical	Augmentation isn't running
Abstract art	Too aggressive
Black frames or wild colors	Normalization problem

06 — Putting It All Together

The final pipeline

One Dataset class: downloads, accesses, transforms, splits, and handles errors.

setup.py

import os
import scipy
import tarfile
import urllib.request
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader, random_split, Subset
from torchvision import transforms


def download_dataset():
    """Download Oxford 102 Flowers Dataset"""
    os.makedirs('flower_data', exist_ok=True)
    for url in [
      'https://www.robots.ox.ac.uk/~vgg/data/flowers/102/102flowers.tgz',
      'https://www.robots.ox.ac.uk/~vgg/data/flowers/102/imagelabels.mat']:
        path = f"flower_data/{url.rsplit('/', 1)[-1]}"
        if not os.path.exists(path):
            urllib.request.urlretrieve(url, path)

    if not os.path.exists('flower_data/jpg'):
        with tarfile.open('flower_data/102flowers.tgz') as tar:
            tar.extractall('flower_data')

download_dataset()

The class takes both transforms and a train flag, and holds its own split sizes and seed.

split() uses random_split for the shuffled indices only, then attaches those indices to fresh views that differ in one field: the train flag. Training gets the augmented transform, validation and test get the plain one. All three share one error log.

oxford_flowers_dataset_final.py

class OxfordFlowersDataset(Dataset):

    # Setup where to find images and labels
    def __init__(self, root_dir, train_transform=None, val_transform=None, train=True,
                 val_size=0.15, test_size=0.15, seed=42):
        self.root_dir = root_dir
        self.img_dir = os.path.join(root_dir, 'jpg')

        labels_mat = scipy.io.loadmat(os.path.join(root_dir, 'imagelabels.mat'))
        self.labels = labels_mat['labels'][0] - 1  # labels start from 1 but pytorch expects from 0
        self.train_transform = train_transform
        self.val_transform = val_transform
        self.train = train

        # split configuration
        self.val_size = val_size
        self.test_size = test_size
        self.seed = seed

        # keep track of errors we encounter
        self.error_log = []

    # Count number of total samples
    def __len__(self):
        return len(self.labels)

    # Build train / val / test views with the right transform on each
    def split(self):
        """Split into three Subsets, each pointing at a dataset 
          in the correct mode"""
        n = len(self)
        val_n = int(self.val_size * n)
        test_n = int(self.test_size * n)
        train_n = n - val_n - test_n

        generator = torch.Generator().manual_seed(self.seed)
        train_split, val_split, test_split = random_split(
            self, [train_n, val_n, test_n], generator=generator
        )

        def view(train):
            ds = OxfordFlowersDataset(self.root_dir,
                                      train_transform=self.train_transform,
                                      val_transform=self.val_transform,
                                      train=train,
                                      val_size=self.val_size,
                                      test_size=self.test_size,
                                      seed=self.seed)
            ds.error_log = self.error_log  # one shared log across all splits
            return ds

        return (Subset(view(True),  train_split.indices),
                Subset(view(False), val_split.indices),
                Subset(view(False), test_split.indices))

    # How to get image and label number 'idx'
    def __getitem__(self, idx, _depth=0):
        "Load image with error handling"
        img_path = 'unknown'
        try:
            img_name = f"image_{idx+1:05d}.jpg"
            img_path = os.path.join(self.img_dir, img_name)
            image = Image.open(img_path)

            # Check for corruption
            image.verify()  # verify closes the file
            image = Image.open(img_path)

            # Skip tiny images
            if image.size[0] < 32 or image.size[1] < 32:
                raise ValueError(f"Image too small: {image.size}")

            # Convert grayscale to RGB
            if image.mode != 'RGB':
                image = image.convert('RGB')

            label = int(self.labels[idx])
            transform = self.train_transform if self.train else self.val_transform
            if transform:
                image = transform(image)
            return image, label
        except Exception as e:
            self.error_log.append({
                'index': idx,
                'error': str(e),
                'path': img_path
            })
            print(f'Warning: Skipping corrupted image {idx} : {e}')
            if _depth >= 10:
                raise RuntimeError(f'Too many consecutive bad images near index {idx}') from e
            next_idx = (idx + 1) % len(self)
            return self.__getitem__(next_idx, _depth + 1)

    def get_error_summary(self):
        """Review after training"""
        if not self.error_log:
            print('No Errors encountered! - dataset is clean!')
        else:
            print(f'Encountered {len(self.error_log)} problematic images')
            for error in self.error_log[:5]:
                print(f"Index {error['index']}:{error['error']}")

    # Where does sample 'idx' live on disk
    def get_image_path(self, idx):
        return os.path.join(self.img_dir, f"image_{idx+1:05d}.jpg")

Finally loading the datasets...

build_the_splits.py

dataset = OxfordFlowersDataset('flower_data',
                               train_transform=train_transform,
                               val_transform=val_transform,
                               val_size=0.15,
                               test_size=0.15,
                               seed=42)

train_dataset, val_dataset, test_dataset = dataset.split()

print(f'Training   : {len(train_dataset)}')
print(f'Validation : {len(val_dataset)}')
print(f'Test       : {len(test_dataset)}')

output

Training   : 5733
Validation : 1228
Test       : 1228

8,189 images into 5,733 / 1,228 / 1,228. One sample from each split returns torch.Size([3, 224, 224]) — augmented on the training side, fixed on the other.

07 — Results

Bonus: Inspecting Augmentations

The best way to know augmentations are working is to visually inspect

This code does exactly the same. It also includes reversing normalization and puts the original file first for comparison. One detail: the splits are Subset objects with shuffled indices, so index 0 of the training split is not the first file on disk. Unwrapping the Subset recovers which file it actually is.

visualize_augmentation.py

import math
import matplotlib.pyplot as plt
from torch.utils.data import Subset

MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
STD  = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

def denormalize(img):
    return (img * STD + MEAN).clamp(0, 1)

def visualize_augmentation(dataset, idx=0, num_versions=7):
    """Visualize Augmentations"""
    # unwrap Subset so idx maps back to the file on disk
    base, base_idx = (dataset.dataset, dataset.indices[idx]) if isinstance(dataset, Subset) else (dataset, idx)
    img_path = base.get_image_path(base_idx)
    original = Image.open(img_path).convert('RGB')

    total = num_versions + 1
    cols = 4
    rows = math.ceil(total / cols)
    fig, axes = plt.subplots(rows, cols, figsize=(3 * cols, 3 * rows))
    axes = axes.flatten()

    axes[0].imshow(original)
    axes[0].set_title(f'ORIGINAL\n{original.size[0]}x{original.size[1]}', color='red')
    axes[0].axis('off')

    for i in range(num_versions):
        img, label = dataset[idx]
        img = denormalize(img)
        axes[i + 1].imshow(img.permute(1, 2, 0))
        axes[i + 1].set_title(f'Version{i+1}')
        axes[i + 1].axis('off')

    for ax in axes[total:]:      # blank out unused cells
        ax.axis('off')

    plt.tight_layout()
    plt.show()
    return img_path

$Original flower plus seven augmented versions, first training sample$

Original at 591×500, then seven draws from the training transform. Each one differs. Each one is still obviously the same flower.

Original flower plus seven augmented versions, second training sample

A second sample confirms it. Not identical, not abstract, no black frames. The augmentation is running and the signal survives.

What this pipeline gets right Labels zero-indexed · filenames offset by one · no loading in __init__ · single-value Resize before CenterCrop · Normalize after ToTensor · third split as remainder · shuffle on training only · augmentation on training only · augmentation checked by eye · one batch tested before training.

From a folder of anonymous JPEGs to shuffled, augmented, fault-tolerant batches. The data is ready.

Taming the Oxford 102Flowers Dataset

Same model, different results

What you're given

Three things that can go wrong

The custom Dataset class

__init__

__len__

__getitem__

Transform pipelines

Solving Different size problem with transforms

Solving PIL Image in dataset with ToTensor()

Normalize spreads them out

ToTensor is a one-way bridge

Debug one step at a time

Splitting and batching

Shuffle, and when not to

Batches and epochs

Two expensive mistakes

Augmentation, errors, and a visual check

Augmentation

One bad file shouldn't end the run

Is the augmentation too strong?

The final pipeline

Bonus: Inspecting Augmentations

Taming the Oxford 102
Flowers Dataset

init

len

getitem

Solving PIL Image in dataset with `ToTensor()`