Data & Augmentation — Programuoki

Code

Augmentation and balancing — write them yourself

The holes are filled; splitting and normalising work. Two parts remain — this time empty, with no holes: augmentation and class balancing.

What is already written, so you need not

make_loaders has grown. Most importantly, it already knows two things that are easy to get wrong, so you do not write them:

# Two views of the same files: augmented for training, untouched for validation.
train_ds = datasets.ImageFolder(root, transform=build_transform(augment, mean, std))
val_ds = datasets.ImageFolder(root, transform=build_transform(False, mean, std))

The same files are read twice: the training copy with random changes, the validation copy never. One shared ImageFolder with augmentation would mean the validation accuracy is measured on different pictures every time (see step 7).

The only new places in data.py are the two that now say raise NotImplementedError.

1. Augmentation — `build_transform`

Each time, training pictures are changed a little differently. You do not need new photographs — the same picture is slightly different every epoch.

if augment:
    steps += [
        random horizontal flip,
        random affine change:
            rotation ±10°, shift up to 10 % in each direction,
            scale from 0.9 to 1.1
    ]

In place of raise NotImplementedError write steps += [ ... ] with transforms.RandomHorizontalFlip() and transforms.RandomAffine(...).

Note the order: build_transform first puts Grayscale, then your augmentation, then ToTensor and Normalize. Random changes act on the picture (PIL), while normalisation acts on the numbers.

Think about your categories. A horizontal flip changes nothing for a shoe or a T-shirt. But if your categories are "letter b" and "letter d", a flip turns one into the other — and you have broken the labels yourself. Augmentation must preserve what the network is supposed to recognise.

2. Class balancing — `make_loaders`

If one class has 20 pictures and the others 200, training batches are mostly made of the big classes. One answer is to draw the rare classes more often:

for each training picture:
    weight = 1 / (how many training pictures are in its class)
sampler = WeightedRandomSampler(weights, num_samples=len(train_idx))

In place of raise NotImplementedError write three lines: a Counter to count the classes (already imported), the list of weights, and a WeightedRandomSampler. The result must land in the variable sampler.

It draws with replacement: the same rare photograph may be drawn several times per epoch. That is one of the costs — see step 7.

The test

test_data.py has grown. The new tests check that:

without augmentation a picture loads identically every time;
with augmentation — differently;
validation pictures with augment=True are the same as without it;
without balancing, classes appear as often as they exist ([8, 32, 32]);
with balancing, the rare class (10 pictures) appears more than 85 times in 5 epochs (without balancing — 40; an even split would be ~120).

Run python -m unittest test_data (you already have the holes filled):

.EE.....E..
======================================================================
ERROR: test_augmentation_changes_training_images_every_time (...)
NotImplementedError: write the augmentation steps
ERROR: test_balancing_draws_the_rare_class_far_more_often (...)
NotImplementedError: write the balancing sampler
ERROR: test_validation_images_are_never_augmented (...)
NotImplementedError: write the augmentation steps
----------------------------------------------------------------------
Ran 11 tests in 0.8s

FAILED (errors=3)

(Shortened.) Eight tests pass, three are waiting for your code.

data.py is still shown with its holes in the panel — the panel shows the skeleton, not your file. In your copy they are filled.

When it does not pass

AssertionError: 7 not greater than 85 : alpha drawn only 7 times in 5 epochs — the weights are inverted: you wrote counts[...] rather than 1.0 / counts[...]. Now big classes are drawn even more often.
AssertionError: 40 not greater than 85 — the sampler was built but not assigned to the variable sampler (it stayed None).
ValueError: sampler option is mutually exclusive with shuffle — if you try changing the DataLoader lines. You need not: shuffle=(sampler is None) already handles it.

data.py MODIFY 25 lines changed

Edit this file in ./

"""data.py — your own images, read from folders: data/<class>/*.png."""

from collections import Counter

import torch
from torch.utils.data import DataLoader, Subset, WeightedRandomSampler
from torchvision import datasets, transforms


def class_counts(root: str) -> dict[str, int]:
    """How many images each class folder holds."""
    ds = datasets.ImageFolder(root)
    counts = Counter(ds.targets)
    return {name: counts[i] for i, name in enumerate(ds.classes)}


def build_transform(
    augment: bool, mean: float | None = None, std: float | None = None
) -> transforms.Compose:
    steps: list = [transforms.Grayscale(num_output_channels=1)]
    if augment:  # random changes, applied freshly every time an image is loaded
        # Add two transforms to steps: a RandomHorizontalFlip, and a
        # RandomAffine with degrees=10, translate=(0.1, 0.1), scale=(0.9, 1.1).
        raise NotImplementedError("write the augmentation steps")
    steps.append(transforms.ToTensor())  # PIL image -> [1, H, W] floats in 0..1
    if mean is not None:
        steps.append(transforms.Normalize([mean], [std]))
    return transforms.Compose(steps)


def split_indices(
    targets: list[int], val_frac: float, seed: int
) -> tuple[list[int], list[int]]:
    """A stratified split: every class contributes val_frac of its images."""
    gen = torch.Generator().manual_seed(seed)
    train_idx: list[int] = []
    val_idx: list[int] = []
    for cls in sorted(set(targets)):
        idx = [i for i, t in enumerate(targets) if t == cls]
        order = torch.randperm(len(idx), generator=gen).tolist()
        # HOLE 1: how many images of this class go to validation?
        # val_frac of them, rounded - but never fewer than one.
        n_val = 0
        val_idx += [idx[j] for j in order[:n_val]]
        train_idx += [idx[j] for j in order[n_val:]]
    return train_idx, val_idx


def mean_std(ds: datasets.ImageFolder, idx: list[int]) -> tuple[float, float]:
    """Pixel mean and std over the TRAINING images only."""
    pixels = torch.cat([ds[i][0].flatten() for i in idx])
    return pixels.mean().item(), pixels.std().item()


def make_loaders(
    root: str,
    batch_size: int = 32,
    val_frac: float = 0.2,
    augment: bool = False,
    balance: bool = False,
    seed: int = 42,
) -> tuple[DataLoader, DataLoader, list[str]]:
    """Return (train_loader, val_loader, class_names)."""
    plain = datasets.ImageFolder(root, transform=build_transform(augment=False))
    train_idx, val_idx = split_indices(plain.targets, val_frac, seed)
    # HOLE 2: the mean and std used to normalise. Measure them with
    # mean_std() - on WHICH images?
    mean, std = 0.0, 1.0

    # Two views of the same files: augmented for training, untouched for validation.
    train_ds = datasets.ImageFolder(root, transform=build_transform(augment, mean, std))
    val_ds = datasets.ImageFolder(root, transform=build_transform(False, mean, std))

    sampler = None
    if balance:  # draw rare classes more often
        # Give every training image the weight 1 / (number of training images in
        # its class), then build a WeightedRandomSampler that draws
        # len(train_idx) images per epoch with those weights.
        raise NotImplementedError("write the balancing sampler")

    train_loader = DataLoader(
        Subset(train_ds, train_idx),
        batch_size=batch_size,
        shuffle=(sampler is None),
        sampler=sampler,
    )
    val_loader = DataLoader(Subset(val_ds, val_idx), batch_size=batch_size)
    return train_loader, val_loader, plain.classes

test_data.py MODIFY 42 lines changed