Evaluation & Deployment

Code

Loading and predicting — write it yourself

evaluate.py can already build a matrix from two tensors: the truth and the predictions. Now it needs to get those tensors from a real model. Two functions, no holes this time: empty skeletons that you fill in completely.

Pseudocode

load_model — turns model.pt back into a working network:

ckpt  ← load the file (onto the CPU)
model ← a new ResNet-18 with as many outputs as ckpt has classes, WITHOUT ImageNet weights
fill model with the weights from ckpt["state_dict"]
switch model to evaluation mode
return model, ckpt["classes"], ckpt["size"]

predict_all — runs every batch through the model:

preds, targets ← two empty lists
for each batch (xb, yb) in loader:
    logits  ← model(xb)
    preds   ← append the INDEX of the highest-scoring class (one per image)
    targets ← append yb
return preds and targets, each glued into one long tensor

A couple of things the pseudocode does not say:

Evaluation mode is not a detail. ResNet-18 has BatchNorm layers; in training mode they normalise by the current batch's statistics and change themselves. In the next step you will see what that does to accuracy.
predict_all is already marked @torch.no_grad() — it needs no gradients.
Every image must appear exactly once, in the loader's order. The matrix and the list of mistakes with file names rely on it.

Your job

evaluate.py (on the right) has grown: load_model and predict_all skeletons with raise NotImplementedError were added, plus image_paths and main, which are ready-made glue (you only read them).

confusion_matrix and class_scores still appear with holes in the panel — the panel shows the skeleton, not your file. In your copy they are already filled in, and you leave them that way.

test_evaluate.py has grown, too. It builds a tiny model (build_resnet(3, pretrained=False) — no internet needed), saves it to a temporary file, and checks that load_model brings it back: the classes, the size, evaluation mode, the number of head outputs, and the same weights.

Run before you write

$ python -m unittest
...
ERROR: test_returns_the_classes_and_the_size_that_were_saved
NotImplementedError
ERROR: test_the_head_has_one_output_per_class
NotImplementedError
ERROR: test_the_model_is_in_evaluation_mode
NotImplementedError
ERROR: test_the_weights_are_the_saved_ones
NotImplementedError
ERROR: test_matches_argmax_in_loader_order
NotImplementedError
Ran 16 tests in 1.116s
FAILED (errors=5)

The eleven old tests pass (they do not depend on the new functions); the five new ones are waiting for you.

When it does not pass

AssertionError: True is not false in the …evaluation_mode test — you did not switch the model to eval(). model.training is still True.
RuntimeError: Error(s) in loading state_dict for ResNet with a line size mismatch for fc.weight: copying a param with shape torch.Size([3, 512]) from checkpoint, the shape in current model is torch.Size([10, 512]) — you rebuilt the model with the wrong number of classes. Take the number from ckpt["classes"], not from your head.
AttributeError: 'list' object has no attribute 'shape' — you returned a list of tensors, not one tensor. Glue them together (torch.cat).
AssertionError: Tuples differ: (10, 3) != (10,) — preds holds the score of every class (3 numbers per image), but you need one number: the class index. Pick the highest score along the class axis (argmax(dim=1)).

evaluate.py MODIFY 72 lines changed

Edit this file in ./

"""evaluate.py — look past one accuracy number: where does the model go wrong?"""

import argparse

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

from data import make_eval_loader, make_loaders
from transfer import build_resnet


def confusion_matrix(
    targets: torch.Tensor, preds: torch.Tensor, n: int
) -> torch.Tensor:
    """cm[i, j] = how many images of true class i were predicted as class j."""
    cm = torch.zeros(n, n, dtype=torch.long)
    for t, p in zip(targets, preds):
        # HOLE 1: count this image. One more image whose TRUE class is t
        # and whose PREDICTED class is p. Rows are the truth.
        pass
    return cm


def class_scores(cm: torch.Tensor, i: int) -> tuple[float, float, float]:
    """(precision, recall, f1) of class i, read off the confusion matrix.

    A class that was never predicted has precision 0.0 - never divide by zero.
    """
    tp = cm[i, i].item()  # true positives: really class i, called class i
    # HOLE 2: precision = tp / (everything PREDICTED as i)
    #         recall    = tp / (everything that REALLY IS i)
    precision = 0.0
    recall = 0.0
    f1 = 2 * precision * recall / max(precision + recall, 1e-9)
    return precision, recall, f1


def print_report(cm: torch.Tensor, classes: list[str]) -> None:
    print(f"{'':<10}" + "".join(f"{c[:8]:>9}" for c in classes) + "   <- predicted")
    for i, c in enumerate(classes):
        print(f"{c:<10}" + "".join(f"{v.item():>9}" for v in cm[i]))

    print(f"\n{'class':<10}{'precision':>10}{'recall':>9}{'f1':>7}{'images':>8}")
    for i, c in enumerate(classes):
        precision, recall, f1 = class_scores(cm, i)
        print(f"{c:<10}{precision:>10.2f}{recall:>9.2f}{f1:>7.2f}{cm[i].sum().item():>8}")


def load_model(path: str) -> tuple[torch.nn.Module, list[str], int]:
    """Rebuild the network from a checkpoint written by transfer.py.

    The checkpoint is a dict: {"state_dict": ..., "classes": [...], "size": int}.
    Return (model, classes, size). The model must:
      - have the SAME architecture, with the right number of output classes;
      - NOT download the ImageNet weights again (they are about to be overwritten);
      - carry the saved weights;
      - be in evaluation mode before it is returned.
    """
    raise NotImplementedError


@torch.no_grad()
def predict_all(
    model: torch.nn.Module, loader: DataLoader
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run every batch of the loader through the model.

    Return (preds, targets): two 1-D tensors, one entry per image, in loader
    order. preds holds the class INDEX with the highest score.
    """
    raise NotImplementedError


def image_paths(loader: DataLoader) -> list[str]:
    ds = loader.dataset
    if isinstance(ds, Subset):
        return [ds.dataset.samples[i][0] for i in ds.indices]
    return [path for path, _ in ds.samples]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="model.pt")
    parser.add_argument("--data", default="data", help="training folder (for the val split)")
    parser.add_argument("--test-dir", default=None, help="evaluate a separate folder instead")
    args = parser.parse_args()

    model, classes, size = load_model(args.model)
    if args.test_dir:
        loader, test_classes = make_eval_loader(args.test_dir, size=size)
        assert test_classes == classes, f"class mismatch: {test_classes} vs {classes}"
        print(f"evaluating every image in {args.test_dir}")
    else:
        _, loader, _ = make_loaders(args.data, size=size, rgb=True)
        print(f"evaluating the validation split of {args.data}")

    preds, targets = predict_all(model, loader)
    acc = (preds == targets).float().mean().item()
    print(f"accuracy: {acc:.1%} on {len(targets)} images\n")

    cm = confusion_matrix(targets, preds, len(classes))
    print_report(cm, classes)

    paths = image_paths(loader)
    wrong = (preds != targets).nonzero().flatten().tolist()
    print(f"\n{len(wrong)} mistakes, first 8:")
    for i in wrong[:8]:
        print(f"  {paths[i]}: truth {classes[targets[i]]}, predicted {classes[preds[i]]}")


if __name__ == "__main__":
    main()

test_evaluate.py MODIFY 52 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 torch import nn
from torch.utils.data import DataLoader, TensorDataset

from data import make_eval_loader
from evaluate import class_scores, confusion_matrix, load_model, predict_all, 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)


class TestLoadModel(unittest.TestCase):
    """load_model: a checkpoint written by transfer.py comes back usable."""

    def setUp(self):
        torch.manual_seed(0)
        self.tmp = tempfile.TemporaryDirectory()
        self.path = os.path.join(self.tmp.name, "tiny.pt")
        self.saved = build_resnet(3, pretrained=False)
        torch.save(
            {"state_dict": self.saved.state_dict(), "classes": ["a", "b", "c"], "size": 64},
            self.path,
        )

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

    def test_returns_the_classes_and_the_size_that_were_saved(self):
        _, classes, size = load_model(self.path)
        self.assertEqual(classes, ["a", "b", "c"])
        self.assertEqual(size, 64)

    def test_the_model_is_in_evaluation_mode(self):
        model, _, _ = load_model(self.path)
        self.assertFalse(model.training)

    def test_the_head_has_one_output_per_class(self):
        model, _, _ = load_model(self.path)
        self.assertEqual(tuple(model(torch.zeros(1, 3, 64, 64)).shape), (1, 3))

    def test_the_weights_are_the_saved_ones(self):
        model, _, _ = load_model(self.path)
        self.assertTrue(torch.equal(model.fc.weight, self.saved.fc.weight))


class TestPredictAll(unittest.TestCase):
    def test_matches_argmax_in_loader_order(self):
        torch.manual_seed(0)
        model = nn.Sequential(nn.Flatten(), nn.Linear(12, 3))
        x = torch.randn(10, 3, 2, 2)
        y = torch.arange(10) % 3
        loader = DataLoader(TensorDataset(x, y), batch_size=4)  # last batch has 2
        preds, targets = predict_all(model, loader)
        with torch.no_grad():
            expected = model(x).argmax(dim=1)
        self.assertEqual(tuple(preds.shape), (10,))
        self.assertTrue(torch.equal(preds, expected))
        self.assertTrue(torch.equal(targets, y))


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