Evaluation & Deployment

Code

A model you can reload, and new images

You know how to build the matrix. But it only becomes useful once you have something to put into it: a model that can be loaded from disk, and images nobody has seen. This step makes two small changes — both written for you, because they are plumbing, not an exercise.

1. `data.py`: a folder of "new images"

make_loaders from lesson 9 splits a folder into training and validation. A fresh test/ folder needs no splitting: it needs every image, untouched, in the same order.

def make_eval_loader(
    root: str, batch_size: int = 32, size: int = 224, rgb: bool = True
) -> tuple[DataLoader, list[str]]:

It calls the same build_transform (same size, same normalisation, no augmentation) and a DataLoader without shuffling. The only new thing is that nothing is split.

2. `transfer.py`: a model you can rebuild

So far transfer.py saved nothing after training — the model vanished with the process. Three changes:

build_resnet(..., pretrained=True) — a new parameter. When you turn it off, a fresh ResNet-18 is created without ImageNet weights and nothing is downloaded. You will need it when we rebuild the architecture only to fill it with saved weights.
--out model.pt — where to save.
torch.save({...}) at the end.

Look at what we save:

torch.save(
    {"state_dict": model.state_dict(), "classes": classes, "size": args.size},
    args.out,
)

Not just the weights. The class list and the size the model was trained with travel along. If you kept only the state_dict you would lose two things:

The class names. ImageFolder sorts classes in alphabetical order, not in the order you had in your head. With a hand-typed list the same picture gets another name:

right class list : ['coat', 'pullover', 'sneaker', 'tshirt'] -> sneaker
a list typed by hand: ['tshirt', 'sneaker', 'pullover', 'coat'] -> pullover

The same image — a sneaker. The model did not err, but every prediction carries the wrong name, and nothing reports it.

The size. A prediction for one image must prepare it the same way as in training: the same size. If you do not save it, you have to guess.

A bare state_dict (122 keys for ResNet-18) has no class list in it.

The tests grew

test_evaluate.py gained two new blocks: make_eval_loader (every image, order unchanged, three channels, the right size) and build_resnet(pretrained=False) (works offline, the head has as many outputs as there are classes). The earlier seven tests are untouched.

$ python -m unittest
...........
Ran 11 tests in 0.361s

OK

Training now saves the model

If you do not have a model.pt yet, train — the same command as in lesson 10, with --finetune, which now writes the file:

python transfer.py --data data --epochs 5 --lr 0.001 --finetune

It takes a few minutes (on our CPU, in two runs — 148 and 199 seconds). model.pt comes out at ~45 MB (ours: 44,785,163 bytes): ResNet-18 with 11.2 million weights.

data.py and transfer.py show the full files in the panel with the changed lines highlighted. The rest stays exactly as you had it in lesson 10.

data.py MODIFY 11 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

# What the pretrained networks were trained with (ImageNet pixel statistics).
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]


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: list[float] | None = None,
    std: list[float] | None = None,
    size: int = 28,
    rgb: bool = False,
) -> transforms.Compose:
    if rgb:  # pretrained nets want 3 channels at a larger size
        steps: list = [transforms.Resize((size, size))]
    else:
        steps = [transforms.Grayscale(num_output_channels=1)]
    if augment:  # random changes, applied freshly every time an image is loaded
        steps += [
            transforms.RandomHorizontalFlip(),
            transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.9, 1.1)),
        ]
    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()
        n_val = max(1, round(len(idx) * val_frac))
        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,
    size: int = 28,
    rgb: bool = False,
) -> 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)
    if rgb:
        mean, std = IMAGENET_MEAN, IMAGENET_STD
    else:
        m, s = mean_std(plain, train_idx)
        mean, std = [m], [s]

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

    sampler = None
    if balance:  # draw rare classes more often
        counts = Counter(plain.targets[i] for i in train_idx)
        weights = [1.0 / counts[plain.targets[i]] for i in train_idx]
        sampler = WeightedRandomSampler(weights, num_samples=len(train_idx))

    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


def make_eval_loader(
    root: str, batch_size: int = 32, size: int = 224, rgb: bool = True
) -> tuple[DataLoader, list[str]]:
    """Every image in a folder, untouched: for a brand-new test folder."""
    mean, std = (IMAGENET_MEAN, IMAGENET_STD) if rgb else ([0.5], [0.5])
    ds = datasets.ImageFolder(
        root, transform=build_transform(False, mean, std, size, rgb)
    )
    return DataLoader(ds, batch_size=batch_size), ds.classes

transfer.py MODIFY 14 lines changed

Edit this file in ./

"""transfer.py — stand on a pretrained ResNet-18 and teach it YOUR categories."""

import argparse
import time

import torch
from torch import nn
from torchvision.models import ResNet18_Weights, resnet18

from data import make_loaders
from train import evaluate, fit


def build_resnet(
    num_classes: int, finetune_last_block: bool = False, pretrained: bool = True
) -> nn.Module:
    # pretrained=True downloads the ImageNet weights once (~45 MB).
    weights = ResNet18_Weights.DEFAULT if pretrained else None
    model = resnet18(weights=weights)

    for p in model.parameters():  # freeze everything it already knows
        p.requires_grad = False
    if finetune_last_block:  # optionally let the last residual block adapt
        for p in model.layer4.parameters():
            p.requires_grad = True

    # A brand-new head, trainable by default: 512 features -> your classes.
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    return model


def count_trainable(model: nn.Module) -> tuple[int, int]:
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return trainable, total


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", default="data", help="folder of class folders")
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--size", type=int, default=224)
    parser.add_argument("--finetune", action="store_true", help="also train layer4")
    parser.add_argument("--augment", action="store_true")
    parser.add_argument("--out", default="model.pt", help="where to save the model")
    args = parser.parse_args()

    torch.manual_seed(42)
    train_loader, val_loader, classes = make_loaders(
        args.data, batch_size=32, augment=args.augment, size=args.size, rgb=True
    )

    model = build_resnet(len(classes), args.finetune)
    trainable, total = count_trainable(model)
    print(f"classes: {classes}")
    print(f"trainable parameters: {trainable:,} of {total:,}")

    _, acc = evaluate(model, val_loader)
    print(f"before training: val accuracy {acc:.1%}  (new head is random)")

    start = time.time()
    fit(model, train_loader, args.epochs, args.lr, val_loader)
    _, val_acc = evaluate(model, val_loader)
    print(f"final: val accuracy {val_acc:.1%}")

    # Weights alone are not enough: keep what predict.py needs to rebuild the model.
    torch.save(
        {"state_dict": model.state_dict(), "classes": classes, "size": args.size},
        args.out,
    )
    print(f"saved {args.out}")
    print(f"done in {time.time() - start:.1f}s")


if __name__ == "__main__":
    main()

test_evaluate.py MODIFY 42 lines changed

Edit this file in ./

"""test_evaluate.py — the verdict for evaluate.py. Run: python -m unittest -v"""

import io
import os
import tempfile
import unittest
from contextlib import redirect_stdout

import torch
from PIL import Image

from data import make_eval_loader
from evaluate import class_scores, confusion_matrix, print_report
from transfer import build_resnet, count_trainable

# Six images, three classes: 0 = cat, 1 = dog, 2 = fox.
# The model never says "fox" - so class 2 is never predicted.
TARGETS = torch.tensor([0, 0, 0, 1, 1, 2])
PREDS = torch.tensor([0, 0, 1, 1, 0, 1])
CLASSES = ["cat", "dog", "fox"]


class TestConfusionMatrix(unittest.TestCase):
    def test_rows_are_the_truth_columns_the_prediction(self):
        cm = confusion_matrix(TARGETS, PREDS, 3)
        self.assertEqual(cm.tolist(), [[2, 1, 0], [1, 1, 0], [0, 1, 0]])

    def test_every_image_is_counted_once(self):
        cm = confusion_matrix(TARGETS, PREDS, 3)
        self.assertEqual(cm.sum().item(), 6)


class TestClassScores(unittest.TestCase):
    def setUp(self):
        self.cm = confusion_matrix(TARGETS, PREDS, 3)

    def test_precision_is_a_column(self):
        # dog was predicted 3 times (1 right) -> 1/3
        precision, _, _ = class_scores(self.cm, 1)
        self.assertAlmostEqual(precision, 1 / 3, places=4)

    def test_recall_is_a_row(self):
        # there are 2 real dogs (1 found) -> 1/2
        _, recall, _ = class_scores(self.cm, 1)
        self.assertAlmostEqual(recall, 1 / 2, places=4)

    def test_f1_combines_both(self):
        _, _, f1 = class_scores(self.cm, 1)
        self.assertAlmostEqual(f1, 0.4, places=4)

    def test_a_class_that_is_never_predicted_does_not_crash(self):
        self.assertEqual(class_scores(self.cm, 2), (0.0, 0.0, 0.0))


class TestReport(unittest.TestCase):
    def test_report_prints_the_scores_per_class(self):
        cm = confusion_matrix(TARGETS, PREDS, 3)
        out = io.StringIO()
        with redirect_stdout(out):
            print_report(cm, CLASSES)
        lines = out.getvalue().splitlines()
        table = lines[next(i for i, l in enumerate(lines) if l.startswith("class")):]
        dog_line = next(l for l in table if l.startswith("dog"))
        self.assertIn("0.33", dog_line)  # precision
        self.assertIn("0.50", dog_line)  # recall


class TestEvalLoader(unittest.TestCase):
    """make_eval_loader: every image in a folder, same preprocessing, no shuffle."""

    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        for cls in ("ant", "bee"):  # 3 grey 28x28 images per class
            os.makedirs(os.path.join(self.tmp.name, cls))
            for k in range(3):
                img = Image.new("L", (28, 28), color=40 * k + 20)
                img.save(os.path.join(self.tmp.name, cls, f"{k}.png"))

    def tearDown(self):
        self.tmp.cleanup()

    def test_batches_are_rgb_and_resized(self):
        loader, classes = make_eval_loader(self.tmp.name, batch_size=4, size=64)
        xb, yb = next(iter(loader))
        self.assertEqual(classes, ["ant", "bee"])
        self.assertEqual(tuple(xb.shape), (4, 3, 64, 64))

    def test_order_is_fixed_and_follows_the_folders(self):
        loader, _ = make_eval_loader(self.tmp.name, batch_size=4, size=64)
        labels = [y for _, yb in loader for y in yb.tolist()]
        self.assertEqual(labels, [0, 0, 0, 1, 1, 1])


class TestBuildResnet(unittest.TestCase):
    def test_pretrained_false_builds_offline_with_the_right_head(self):
        model = build_resnet(3, pretrained=False)
        self.assertEqual(model.fc.out_features, 3)

    def test_only_the_new_head_is_trainable_by_default(self):
        model = build_resnet(3, pretrained=False)
        trainable, _ = count_trainable(model)
        self.assertEqual(trainable, 512 * 3 + 3)


if __name__ == "__main__":
    unittest.main()

A model you can reload, and new images

1. data.py: a folder of "new images"

2. transfer.py: a model you can rebuild

The tests grew

Training now saves the model

1. `data.py`: a folder of "new images"

2. `transfer.py`: a model you can rebuild