Hands-On: Fine-Tune a ResNet on Your Own Photos

What you will build

In part 6 of this path you learned the idea behind transfer learning: take a network that already knows how to see, such as a ResNet trained on millions of ImageNet photos, and reuse its convolutional layers for a new, much smaller problem. In part 4 you saw why ResNet in particular became a default choice, thanks to residual connections that let very deep networks train without their gradients vanishing. In part 11 you learned how to look inside a trained CNN with Grad-CAM to check whether it is paying attention to the right things.

This article puts all of that to work. You will organize a small image dataset the way PyTorch expects it, load a pretrained ResNet18, replace its final layer, fine-tune it in two phases, evaluate it properly, inspect its mistakes, and visualize what it looks at. The code below uses a small dataset of generated images so that everything in this article runs end to end without you needing to download anything first. Once it works, swapping in your own photos is a matter of replacing one folder of images with another, which the final section walks through.

The code blocks are meant to be run in order, in the same notebook or script, because later blocks reuse variables, models and imports created earlier. Each block is complete on its own in the sense that it will run correctly at that point in the notebook; it is not meant to be copy-pasted in isolation out of sequence.

Before you start

This tutorial assumes Python 3.10 or later, PyTorch 2.x, and torchvision 0.17 or later. The torchvision version matters because it introduces the weights=ResNet18_Weights.DEFAULT style of loading pretrained models, which replaced the older pretrained=True argument. If you have an older torchvision, the rest of the code still works conceptually, but you would load the model with resnet18(pretrained=True) instead of the weights enum shown below. You will also need matplotlib for plotting and scikit-learn for the evaluation metrics.

pip install torch torchvision matplotlib scikit-learn pillow

If you have a GPU available, PyTorch will use it automatically through the device selection code below. If you are on an Apple Silicon Mac, PyTorch can also use the Metal backend (mps). Everything in this article will still run on a plain CPU, just more slowly, because the dataset here is intentionally small.

Step 1: organize your photos like a dataset

Almost every image classification tool in PyTorch's ecosystem expects images arranged in a specific folder layout called an image folder dataset. Each split (train, validation, test) is a top-level folder, and inside each split, every class gets its own subfolder containing only images of that class. For a project classifying three kinds of kitchen tools, it would look like this:

data/
  train/
    mug/
      mug_001.jpg
      mug_002.jpg
      ...
    bottle/
      bottle_001.jpg
      ...
    spoon/
      spoon_001.jpg
      ...
  val/
    mug/
    bottle/
    spoon/
  test/
    mug/
    bottle/
    spoon/

torchvision.datasets.ImageFolder reads this layout automatically: it treats each subfolder name as a class label and builds the label list for you, in alphabetical order of the folder names. This is convenient, but it also means you should never mix images from different classes into the same folder, and you should keep the split folders separate so that no single photo appears in both train and validation. A validation set that accidentally contains near-duplicates of training images will report accuracy that is too optimistic, a problem usually called data leakage.

For your own photos, a reasonable starting point is at least 50 to 100 images per class for training, plus 15 to 25 per class for validation and a similar amount held back for a final test set you only look at once. Fewer than that is still workable with transfer learning, because the pretrained convolutional layers already know how to detect edges, textures and shapes; you are only asking the final layer to learn a new mapping from those features to your classes. But the fewer photos you have, the more careful you need to be about variety: different backgrounds, different lighting, different angles, so the network does not quietly learn to recognize your kitchen counter instead of your mug.

Step 2: a stand-in dataset you can run right now

To keep this tutorial self-contained, we will generate a small synthetic dataset of three simple shapes, circles, squares and triangles, drawn on noisy backgrounds. It plays the same role your own photos would play: three folders of images with the same ImageFolder layout as above. The point of fine-tuning a pretrained ResNet on shapes is not realistic by itself, but every step of the pipeline, loading, augmenting, training, evaluating, and visualizing, is identical to what you would do with real photos, and the small dataset trains in well under a minute even on a CPU. Once it works, swap the folder for your own photos and nothing else in the code needs to change beyond the number of classes.

import os
import random
from pathlib import Path
from PIL import Image, ImageDraw

random.seed(0)

def make_shape_image(shape, size=128):
    bg = (random.randint(180, 255), random.randint(180, 255), random.randint(180, 255))
    img = Image.new("RGB", (size, size), color=bg)
    draw = ImageDraw.Draw(img)
    cx = size // 2 + random.randint(-15, 15)
    cy = size // 2 + random.randint(-15, 15)
    r = random.randint(28, 46)
    color = (random.randint(0, 120), random.randint(0, 120), random.randint(0, 120))
    if shape == "circle":
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill=color)
    elif shape == "square":
        draw.rectangle([cx - r, cy - r, cx + r, cy + r], fill=color)
    elif shape == "triangle":
        draw.polygon([(cx, cy - r), (cx - r, cy + r), (cx + r, cy + r)], fill=color)
    # a little salt-and-pepper noise so the task is not trivially easy
    for _ in range(200):
        x, y = random.randint(0, size - 1), random.randint(0, size - 1)
        img.putpixel((x, y), (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    return img

def build_dataset(root="data", classes=("circle", "square", "triangle"),
                   n_train=80, n_val=20, n_test=20):
    for split, n in (("train", n_train), ("val", n_val), ("test", n_test)):
        for cls in classes:
            folder = Path(root) / split / cls
            folder.mkdir(parents=True, exist_ok=True)
            for i in range(n):
                img = make_shape_image(cls)
                img.save(folder / f"{cls}_{i}.png")

build_dataset()
print("Dataset created under ./data")

Running this creates data/train, data/val and data/test, each with circle, square and triangle subfolders, 80, 20 and 20 images per class respectively. This mirrors exactly what you would have if you had taken photos of three objects on your phone and sorted them into these same folders. When you move to your own photos, you can delete the data folder this code produced and copy your own images in with the same names and layout, and every block from here on will work unchanged, aside from the class names that ImageFolder will pick up automatically.

Step 3: loading images with transforms

Part 5 of this path covered data augmentation: randomly flipping, cropping or recoloring training images so the network sees more variety than the raw dataset contains, which reduces overfitting. We use the same idea here, but more gently than you might for training a CNN from scratch, because we are adapting a network that already has good features, not teaching it to see from zero. Aggressive augmentation on a small fine-tuning dataset can sometimes do more harm than good, since there are fewer images to average the noise out over.

There is one detail that is easy to miss and important to get right: the pretrained ResNet was trained on ImageNet images that were normalized using ImageNet's own per-channel mean and standard deviation. If you feed it images normalized differently, or not normalized at all, the pretrained weights will see inputs in a numerical range they were never trained for, and performance will suffer noticeably, often in confusing ways that look like a bug elsewhere. The three numbers below are the standard ImageNet statistics used throughout the torchvision model zoo.

import torch
from torchvision import transforms

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

train_transform = transforms.Compose([
    transforms.Resize((160, 160)),
    transforms.RandomResizedCrop(128, scale=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

eval_transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

The training transform resizes images a little larger than the final size, then takes a random crop within 80 to 100 percent of the image area, which acts like a gentle zoom-and-pan augmentation, followed by a random horizontal flip and mild color jitter. The evaluation transform, used for validation and test images, skips all of that randomness: it resizes directly to the final size and normalizes, so that evaluating the same image twice always gives the same result. This asymmetry matters: you augment training data to make learning harder and more general, but you want evaluation to be deterministic and fair.

If your own photos include things like text or asymmetric logos where a left-right flip would change the meaning of the class, drop RandomHorizontalFlip. If lighting conditions in your photos are already highly varied, you can reduce or remove ColorJitter. Augmentation choices should reflect what kind of variation you expect to see in new, unseen photos, not variation for its own sake.

from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

train_ds = ImageFolder("data/train", transform=train_transform)
val_ds = ImageFolder("data/val", transform=eval_transform)
test_ds = ImageFolder("data/test", transform=eval_transform)

train_loader = DataLoader(train_ds, batch_size=16, shuffle=True, num_workers=2)
val_loader = DataLoader(val_ds, batch_size=16, shuffle=False, num_workers=2)
test_loader = DataLoader(test_ds, batch_size=16, shuffle=False, num_workers=2)

class_names = train_ds.classes
print("Classes:", class_names)
print("Train size:", len(train_ds), "Val size:", len(val_ds), "Test size:", len(test_ds))

ImageFolder has already done the work of scanning the folders and assigning integer labels; class_names stores the mapping back to readable folder names, in alphabetical order, so for our shapes it will print ['circle', 'square', 'triangle']. The DataLoader wraps each dataset so that training code can pull batches of images and labels at once instead of one at a time, and shuffle=True on the training loader means each epoch sees the images in a different random order, which prevents the network from learning anything about the order itself. If num_workers=2 causes issues on Windows or inside some notebook environments, set it to 0; it only affects how many background processes load images in parallel, not correctness.

Step 4: sanity-check a batch before training anything

Before spending any time training, it is worth looking at a batch of images with their labels to catch mistakes early: wrong folder, flipped labels, or a transform that distorts the image beyond recognition. This single step catches a surprising number of real-world bugs.

import matplotlib.pyplot as plt

def denormalize(img_tensor):
    mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
    std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
    return img_tensor * std + mean

images, labels = next(iter(train_loader))
fig, axes = plt.subplots(2, 4, figsize=(10, 5))
for ax, img, label in zip(axes.flatten(), images, labels):
    img = denormalize(img).permute(1, 2, 0).clamp(0, 1).numpy()
    ax.imshow(img)
    ax.set_title(class_names[label])
    ax.axis("off")
plt.tight_layout()
plt.show()

Because the training transform normalizes images, we need to undo that before displaying them, which is what denormalize does: it multiplies by the standard deviation and adds back the mean, the exact inverse of transforms.Normalize. Running this should show eight small images, each with a visible shape and a title that correctly names it, with slightly different crops and colors each time you re-run the cell because of the random augmentation. If the titles do not match what you see, double check the folder names and that no images leaked into the wrong class folder.

Step 5: load a pretrained ResNet and replace its head

Recall from part 4 that ResNet is built from residual blocks: each block learns a small correction to add to its input rather than learning a full transformation from scratch, which is what allows ResNet variants to stack far more layers than earlier architectures like VGG without training becoming unstable. ResNet18, the variant we use here, has 18 layers with learnable weights organized into four stages (commonly named layer1 through layer4), followed by global average pooling and a single fully connected layer that produces the final class scores. When this network was trained on ImageNet, that final layer had 1000 outputs, one per ImageNet class. We do not want those 1000 outputs; we want however many classes our own dataset has, so the final layer has to be replaced.

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

device = torch.device(
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
print("Using device:", device)

weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)

for param in model.parameters():
    param.requires_grad = False

num_features = model.fc.in_features
model.fc = nn.Linear(num_features, len(class_names))

model = model.to(device)
print(model.fc)

ResNet18_Weights.DEFAULT points torchvision to the best available pretrained weights for this architecture, downloading them the first time you run this (they are cached afterwards, so later runs are instant). The loop that sets requires_grad = False on every parameter freezes the entire network: PyTorch will still compute predictions through these layers, but it will not compute gradients for them during backpropagation, so their weights never change. This is the feature extraction setup from part 6: the pretrained convolutional layers become a fixed feature extractor, and only a new layer on top gets trained.

model.fc.in_features tells us how many features the layer just before the final classifier produces, which for ResNet18 is 512, regardless of how many output classes you need. We replace model.fc with a brand new nn.Linear(512, len(class_names)) layer. A freshly created layer like this has random weights and, crucially, requires_grad=True by default, so it is the only part of the network that will actually learn in the next step.

Step 6: choose a fine-tuning strategy

Part 6 described two broad strategies for adapting a pretrained network: feature extraction, where you freeze everything except a new head, and fine-tuning, where you also update some or all of the pretrained layers, usually with a smaller learning rate so you do not wreck what the network already knows. In practice, a two-phase approach that combines both tends to work better than either alone, especially on small datasets, and it is what we will do here.

Phase A: train only the new head

Right after replacing model.fc, its weights are random. If you immediately unfreeze the whole network and start training, the large, noisy gradients coming from that random layer flow backwards into the pretrained convolutional layers and can disturb weights that took a very large dataset to learn, a phenomenon sometimes described as catastrophic forgetting. Training only the head first lets it settle into reasonable weights while the backbone stays untouched, so that if you do unfreeze layers afterwards, the gradients flowing back are smaller and less destructive.

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=1e-3)

optim.Adam(model.fc.parameters(), lr=1e-3) tells the optimizer to only look at the parameters of model.fc, which matches what we froze above, although in this case it would not strictly matter even if we passed model.parameters() here, since the frozen ones have requires_grad=False and PyTorch skips updating them regardless. Being explicit about which parameters you are optimizing is still good practice, because it protects you from accidentally training something you meant to keep frozen.

The training and evaluation functions

Before running any epochs, it helps to write two small functions: one that runs a single training epoch and reports the average loss and accuracy, and one that does the same thing on a dataset without updating any weights, which we will use for both validation and test evaluation.

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    running_loss, running_correct, total = 0.0, 0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)
        preds = outputs.argmax(dim=1)
        running_correct += (preds == labels).sum().item()
        total += images.size(0)
    return running_loss / total, running_correct / total


def evaluate(model, loader, criterion, device):
    model.eval()
    running_loss, running_correct, total = 0.0, 0, 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            running_loss += loss.item() * images.size(0)
            preds = outputs.argmax(dim=1)
            running_correct += (preds == labels).sum().item()
            total += images.size(0)
    return running_loss / total, running_correct / total

model.train() and model.eval() switch the network between two modes that behave differently for layers such as batch normalization and dropout. ResNet18 uses batch normalization throughout; in training mode it normalizes each batch using that batch's own statistics and keeps a running estimate for later, while in evaluation mode it uses the running estimate instead, which keeps predictions consistent regardless of what else happens to be in the same batch. Forgetting to call model.eval() before evaluating is a common source of oddly inconsistent validation numbers, so it is worth forming the habit of calling it every time you measure performance without training.

Inside evaluate, torch.no_grad() tells PyTorch not to track operations for backpropagation at all, which both saves memory and speeds things up, since we have no intention of calling backward() during evaluation.

num_epochs_a = 5
history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}

for epoch in range(num_epochs_a):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)
    history["train_loss"].append(train_loss)
    history["train_acc"].append(train_acc)
    history["val_loss"].append(val_loss)
    history["val_acc"].append(val_acc)
    print(f"[Phase A] epoch {epoch+1}/{num_epochs_a} "
          f"train_loss={train_loss:.3f} train_acc={train_acc:.3f} "
          f"val_loss={val_loss:.3f} val_acc={val_acc:.3f}")

Exact numbers will vary between runs due to random initialization and augmentation, so do not expect to reproduce any specific figures here. What you should see on this easy synthetic shapes task is training and validation accuracy both climbing quickly over these five epochs, since the pretrained features already separate circles, squares and triangles well and the head only has to learn a simple mapping on top of them. On real photos with harder, more ambiguous classes, expect slower and noisier improvement, and a validation accuracy that sits meaningfully below training accuracy once overfitting starts to creep in.

Phase B: unfreeze part of the backbone

With the head trained, we can now unfreeze some of the backbone, typically the last residual stage, layer4, since the earliest layers tend to learn very general features such as edges and simple textures that transfer well to almost any image task, while the later layers encode more task-specific combinations that benefit most from adapting to your particular classes. We give the unfrozen backbone layers a smaller learning rate than the head, because they start from a good place already and only need gentle nudging, while the head can keep learning at a faster pace. This is the idea of differential, or discriminative, learning rates.

for param in model.layer4.parameters():
    param.requires_grad = True

optimizer = optim.Adam([
    {"params": model.fc.parameters(), "lr": 1e-3},
    {"params": model.layer4.parameters(), "lr": 1e-4},
])

scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

Adam here receives a list of parameter groups instead of a single list of parameters, each with its own learning rate: the head trains at 1e-3, ten times faster than layer4's 1e-4. The scheduler reduces every group's learning rate by a factor of 10 every three epochs, which lets training take larger steps early on and settle into smaller, more careful steps later, a common pattern for avoiding overshooting a good solution once you are close to it.

num_epochs_b = 5

for epoch in range(num_epochs_b):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)
    scheduler.step()
    history["train_loss"].append(train_loss)
    history["train_acc"].append(train_acc)
    history["val_loss"].append(val_loss)
    history["val_acc"].append(val_acc)
    print(f"[Phase B] epoch {epoch+1}/{num_epochs_b} "
          f"train_loss={train_loss:.3f} train_acc={train_acc:.3f} "
          f"val_loss={val_loss:.3f} val_acc={val_acc:.3f}")

We append to the same history dictionary so the full ten-epoch run, five from each phase, can be plotted together. On the synthetic dataset you will likely see validation accuracy already close to its ceiling by the end of Phase A, with Phase B offering only a small additional gain, if any, simply because the task is easy and the pretrained features were already enough. On a harder, more realistic dataset, Phase B is often where the larger jumps in accuracy happen, because now the network can adjust higher-level features to your specific classes rather than relying on generic ImageNet features alone.

plt.figure(figsize=(10, 4))

plt.subplot(1, 2, 1)
plt.plot(history["train_loss"], label="train")
plt.plot(history["val_loss"], label="val")
plt.axvline(x=num_epochs_a - 0.5, color="gray", linestyle="--", label="phase boundary")
plt.xlabel("epoch")
plt.ylabel("loss")
plt.legend()
plt.title("Loss")

plt.subplot(1, 2, 2)
plt.plot(history["train_acc"], label="train")
plt.plot(history["val_acc"], label="val")
plt.axvline(x=num_epochs_a - 0.5, color="gray", linestyle="--", label="phase boundary")
plt.xlabel("epoch")
plt.ylabel("accuracy")
plt.legend()
plt.title("Accuracy")

plt.tight_layout()
plt.show()

The dashed vertical line marks where Phase A ends and Phase B begins. A healthy pattern looks like both loss curves decreasing and both accuracy curves increasing, with training and validation staying reasonably close together. If training loss keeps falling while validation loss starts rising, that is the classic signature of overfitting: the network is starting to memorize details of the training images rather than learning features that generalize. We will come back to what to do about that in the troubleshooting section.

Step 8: evaluate on the held-out test set

The validation set was used during training to monitor progress, which means your choices, such as how many epochs to run or which layers to unfreeze, were implicitly influenced by it. The test set, which the training loop never touched, gives a cleaner estimate of how the model would perform on genuinely new photos. It is good practice to look at the test set exactly once, at the end, rather than repeatedly tuning against it.

from sklearn.metrics import classification_report, confusion_matrix

model.eval()
all_preds, all_labels = [], []
with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        outputs = model(images)
        preds = outputs.argmax(dim=1).cpu().numpy()
        all_preds.extend(preds)
        all_labels.extend(labels.numpy())

print(classification_report(all_labels, all_preds, target_names=class_names))
print("Confusion matrix:")
print(confusion_matrix(all_labels, all_preds))

classification_report prints precision, recall and F1-score per class, plus overall accuracy. Precision for a class answers: of everything the model called this class, how much actually was; recall answers: of everything that actually was this class, how much did the model catch. These two numbers can disagree sharply when classes are imbalanced, which is why looking at both, rather than only overall accuracy, matters once you move beyond a toy balanced dataset like this one.

The confusion matrix is a grid where rows are true classes and columns are predicted classes, in the order given by class_names; the entry in row i, column j counts how many images of true class i were predicted as class j. A perfect classifier would have all its mass on the diagonal. For our shapes task, expect something close to a clean diagonal; any confusion that does appear is worth investigating with the next step.

Step 9: look at what the model gets wrong

Aggregate numbers hide individual failures. Looking directly at misclassified images is often the fastest way to understand whether a model's mistakes are reasonable (two genuinely similar-looking shapes confused for each other) or a sign of a deeper problem (a bug in labels, or a model keying off something irrelevant like background color).

def collect_misclassified(model, dataset, device, max_items=8):
    model.eval()
    loader = DataLoader(dataset, batch_size=1, shuffle=False)
    mistakes = []
    with torch.no_grad():
        for idx, (image, label) in enumerate(loader):
            image = image.to(device)
            output = model(image)
            pred = output.argmax(dim=1).item()
            if pred != label.item():
                mistakes.append((idx, pred, label.item()))
            if len(mistakes) >= max_items:
                break
    return mistakes

mistakes = collect_misclassified(model, test_ds, device)
for idx, pred, true in mistakes:
    print(f"test image #{idx}: predicted '{class_names[pred]}', actually '{class_names[true]}'")

This walks through the test set one image at a time, which is slower than batching but simplifies tracking which exact image index each mistake corresponds to, so you can go back and look at that file directly. If this list comes back empty on the shapes dataset, that just means the model got every test image right, which is plausible on a task this easy; on real photos with genuinely ambiguous cases, you should expect to see at least a handful of mistakes here even from a well-trained model.

fig, axes = plt.subplots(1, max(len(mistakes), 1), figsize=(3 * max(len(mistakes), 1), 3))
if len(mistakes) == 1:
    axes = [axes]
for ax, (idx, pred, true) in zip(axes, mistakes):
    image, _ = test_ds[idx]
    img = denormalize(image).permute(1, 2, 0).clamp(0, 1).numpy()
    ax.imshow(img)
    ax.set_title(f"pred: {class_names[pred]}\ntrue: {class_names[true]}", fontsize=9)
    ax.axis("off")
plt.tight_layout()
plt.show()

If you see a pattern in the mistakes, for example a model that confuses two particular classes much more than others, it is worth going back to Step 1 and asking whether those two classes have enough visual difference in your photos, or whether you need more examples that show the distinguishing detail clearly.

Step 10: see what the network is looking at

Part 11 introduced Grad-CAM: a technique that uses the gradients flowing into the last convolutional layer to produce a coarse heatmap showing which regions of an image most influenced a particular prediction. We now implement it directly against our fine-tuned model, using PyTorch hooks, functions you register on a layer that get called automatically during the forward and backward passes.

import torch.nn.functional as F

class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.gradients = None
        self.activations = None
        target_layer.register_forward_hook(self._save_activation)
        target_layer.register_full_backward_hook(self._save_gradient)

    def _save_activation(self, module, input, output):
        self.activations = output.detach()

    def _save_gradient(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()

    def __call__(self, input_tensor, class_idx=None):
        self.model.zero_grad()
        output = self.model(input_tensor)
        if class_idx is None:
            class_idx = output.argmax(dim=1).item()
        score = output[0, class_idx]
        score.backward()
        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
        cam = (weights * self.activations).sum(dim=1, keepdim=True)
        cam = F.relu(cam)
        cam = F.interpolate(cam, size=input_tensor.shape[2:], mode="bilinear", align_corners=False)
        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
        return cam, class_idx

gradcam = GradCAM(model, model.layer4[-1])

register_forward_hook attaches _save_activation to model.layer4[-1], the last block of the final residual stage, so that every time an image passes through the network, we capture that layer's output feature maps. register_full_backward_hook attaches _save_gradient so that when we call backward() on a chosen class score, we capture the gradient of that score with respect to the same feature maps. Grad-CAM's core idea, as covered conceptually in part 11, is that averaging those gradients over space gives a weight per feature channel describing how important that channel was for this prediction; multiplying the feature maps by those weights and summing over channels produces a single map that highlights the regions the network relied on most. The ReLU keeps only the positive contributions, since negative values would indicate regions that pushed the score down, which is not what we want to highlight as supporting evidence. Resizing and normalizing to the 0 to 1 range just makes the result easy to overlay on the original image.

def show_gradcam(model, gradcam, dataset, idx, device, class_names):
    image, label = dataset[idx]
    input_tensor = image.unsqueeze(0).to(device)
    cam, pred_idx = gradcam(input_tensor)

    img = denormalize(image).permute(1, 2, 0).clamp(0, 1).numpy()

    fig, axes = plt.subplots(1, 2, figsize=(7, 3.5))
    axes[0].imshow(img)
    axes[0].set_title(f"true: {class_names[label]}")
    axes[0].axis("off")

    axes[1].imshow(img)
    axes[1].imshow(cam, cmap="jet", alpha=0.45)
    axes[1].set_title(f"predicted: {class_names[pred_idx]}")
    axes[1].axis("off")

    plt.tight_layout()
    plt.show()

show_gradcam(model, gradcam, test_ds, idx=0, device=device, class_names=class_names)

The left panel shows the plain test image, the right panel overlays the Grad-CAM heatmap in red and yellow over the areas that contributed most to the predicted class. For our shapes, a sensible result is a hot region centered on the shape itself rather than the noisy background. If you instead saw the heatmap consistently lighting up a corner of the image or the background rather than the object, that would be a warning sign that the model latched onto something incidental, a lesson that applies just as much, if not more, to real photos: a model trained to recognize your dog, for example, might quietly learn to recognize your living room instead if every training photo of the dog happens to be taken there.

Step 11: save the model, reload it, and predict on a new photo

Once you are satisfied with a model's performance, you want to save it so you do not have to retrain every time you restart your notebook or deploy it elsewhere. PyTorch's recommended approach is to save the state dictionary, a plain mapping from layer names to tensors of weights, rather than the whole model object, because state dictionaries are more portable across code changes and PyTorch versions.

torch.save(model.state_dict(), "resnet18_finetuned.pt")
print("Saved model weights to resnet18_finetuned.pt")

To load it back, you need to recreate the same architecture first, then load the saved weights into it. This is why it matters to remember how many classes the model was trained for and what the final layer looked like.

loaded_model = resnet18(weights=None)
loaded_model.fc = nn.Linear(loaded_model.fc.in_features, len(class_names))
loaded_model.load_state_dict(torch.load("resnet18_finetuned.pt", map_location=device))
loaded_model = loaded_model.to(device)
loaded_model.eval()
print("Model reloaded and ready for inference")

weights=None skips downloading the pretrained ImageNet weights, since we are about to overwrite every weight with our own saved ones anyway. map_location=device makes sure the tensors land on whatever device you are using now, even if they were saved from a different device, for example saved on a GPU machine and loaded on a laptop with no GPU.

def predict_image(model, image_path, transform, class_names, device):
    image = Image.open(image_path).convert("RGB")
    tensor = transform(image).unsqueeze(0).to(device)
    model.eval()
    with torch.no_grad():
        output = model(tensor)
        probs = F.softmax(output, dim=1).squeeze().cpu().numpy()
    pred_idx = int(probs.argmax())
    return class_names[pred_idx], float(probs[pred_idx])

sample_path = "data/test/circle/circle_0.png"
label, confidence = predict_image(loaded_model, sample_path, eval_transform, class_names, device)
print(f"Predicted: {label} ({confidence:.1%} confidence)")

convert("RGB") guards against occasional images that are grayscale or have a transparency channel, which would otherwise produce a tensor with the wrong number of color channels and crash the model. softmax turns the raw output scores into probabilities that sum to 1 across classes, and the confidence reported is simply the probability the model assigned to its own top prediction. This same function is what you would wire up to a file upload, a webcam frame, or an API endpoint in a real application; the model itself does not care where the image came from, only that it goes through the same eval_transform used during validation.

Common mistakes and how to avoid them

• Forgetting to normalize with ImageNet statistics, or normalizing with the wrong numbers. The pretrained weights expect inputs in roughly the same distribution they were trained on; skipping or mismatching normalization quietly degrades accuracy without producing an obvious error.
• Calling model.eval() only sometimes. Batch normalization and dropout behave differently in train and eval mode; leaving a model in train mode while reporting validation numbers gives inconsistent, batch-dependent results.
• Unfreezing the whole backbone immediately with a learning rate meant for a freshly initialized head. Large gradients from a random final layer can erase useful pretrained weights in the first few steps; train the head alone first, then unfreeze gradually with a smaller learning rate for the backbone.
• Data leakage between train and validation, for example near-duplicate photos (the same object photographed seconds apart) ending up in both splits. This inflates validation accuracy and hides real generalization problems until you try the model on genuinely new photos.
• Evaluating on the test set repeatedly while tuning hyperparameters. Once you start adjusting your approach based on test results, the test set stops giving an honest estimate of generalization; use the validation set for that and look at the test set only at the end.
• Highly imbalanced classes without any adjustment. If one class has ten times as many photos as another, a model can get decent-looking overall accuracy by mostly ignoring the rare class. Passing class weights to CrossEntropyLoss (criterion = nn.CrossEntropyLoss(weight=class_weights)) or balancing the dataset directly helps here.
• Too few validation images to trust the reported accuracy. An accuracy of 90 percent on a validation set of only 10 images means 1 more mistake would have given 80 percent; small validation sets are noisy, so treat differences of a few percentage points with caution unless your validation set is reasonably sized.
• Resizing images to a very different aspect ratio than they were photographed in, squashing or stretching objects in ways that do not occur in real photos the model will later see.

Adapting this to your own photos

With the pipeline working end to end on generated shapes, moving to your own photos is mostly a matter of replacing the data folder and adjusting a few numbers to match your situation.

1. Delete or rename the data folder created in Step 2, and create a new one with the same train, val, test layout, with your own photos sorted into class subfolders.
2. Re-run the ImageFolder and DataLoader code in Step 3; class_names will update automatically to match your new folder names, and nothing downstream needs to change to account for a different number of classes, since model.fc was sized using len(class_names).
3. If your photos are naturally larger or smaller than typical camera resolution, you can adjust the Resize and RandomResizedCrop sizes in the transforms, though 128 to 224 pixels is a reasonable range for ResNet18 and keeps training fast.
4. Start with Phase A alone for a few more epochs than shown here if your dataset is small, since the head needs enough examples to settle; watch the validation curves rather than a fixed epoch count.
5. Only move to Phase B, unfreezing layer4, once Phase A's validation accuracy has clearly leveled off. If your dataset is very small, a handful of classes with only a few dozen images each, you may find full fine-tuning offers little benefit over feature extraction alone, and risks overfitting instead.
6. If classes are imbalanced, compute class weights from the training set counts and pass them to CrossEntropyLoss, or oversample the rarer classes when building the dataset.
7. Once satisfied, save the model as shown in Step 11, and keep the exact eval_transform used during training alongside it; a model is only as useful as the preprocessing pipeline that feeds it, and mismatched preprocessing at inference time is a common source of silent accuracy drops.

One more practical note: if you want a sense of how much training is happening in each phase without waiting for full training runs, watch the gap between training and validation accuracy as much as the absolute numbers. A small, stable gap suggests the model is learning features that generalize; a widening gap over epochs is the earliest and clearest sign of overfitting, well before validation accuracy itself starts to fall.

Summary and what's next

You took a ResNet18 pretrained on ImageNet, replaced its final layer for a new set of classes, and fine-tuned it in two phases: first training only the new head while the backbone stayed frozen, then unfreezing the last residual stage with a smaller learning rate to adapt higher-level features to your own classes. Along the way you built an image-folder dataset, applied augmentation and ImageNet normalization correctly, wrote reusable training and evaluation functions, checked performance with a classification report and confusion matrix, inspected individual mistakes, used Grad-CAM to see what the network attended to, and saved the model for later use with a simple prediction function. Every one of these steps carries over directly to real photographs; only the data folder changes.

The next article in this path, part 13, moves from classification to detection: instead of asking what is in an image, you will train a YOLO detector to find where objects are, on a custom dataset you label yourself. Part 14 then brings several ideas from across this path together into a complete project, a plant disease classifier served as an API, which reuses the fine-tuning workflow from this article as one piece of a larger, deployable system.