Your First Image Classifier (PyTorch + Transfer Learning)

Everything so far has been features you designed by hand. This post hands that job to the model. You will train a working classifier on 1,200 photos — a dataset far too small to train a network from scratch — by reusing one that has already learned to see.

Prerequisites: The CV project workflow, evaluation metrics, and ideally neural networks.

The problem: sorting recycling photos

A recycling facility wants a phone app that tells people which bin an item goes in. Four classes: cardboard, plastic, glass, metal. A summer intern photographed items over three weeks and collected 1,200 images.

Two problems, both visible before writing a line of model code. Cardboard has 3.7 times more examples than metal. And glass and metal were both shot by intern B in week 3, so anything specific to intern B’s phone or lighting is a shortcut the model can learn instead of learning what glass looks like.

Why 1,200 images is enough

ResNet-18 has 11.7 million parameters. With 1,200 images that is roughly 10,000 parameters per training example. A model that flexible will memorise the training set perfectly and learn nothing that generalises.

The trick is that most of those 11.7 million parameters do not need to change. A network trained on ImageNet’s 1.2 million photos already knows edges, textures, and object parts. Cardboard corrugation is a texture; a bottle is a shape. Those detectors are already there.

Step through the visualization above and watch which layers update. In phase one only the final layer moves. In phase two the later blocks join in, gently.

That 2,052 comes straight from the architecture. ResNet-18 outputs a 512-dimensional feature vector, and a 4-class head is a linear layer:

$512 \times 4 + 4 = 2048 + 4 = 2052 \text{ parameters}$

Fitting 2,052 parameters to 1,200 images is a reasonable ratio. Fitting 11.7 million to 1,200 is not.

graph LR
  A["Input 224×224×3"] --> B["conv1 + pool<br/>FROZEN"]
  B --> C["layer1, layer2<br/>FROZEN"]
  C --> D["layer3<br/>FROZEN"]
  D --> E["layer4<br/>fine-tune, low LR"]
  E --> F["avgpool → 512"]
  F --> G["new Linear 512→4<br/>TRAIN, higher LR"]
  G --> H["4 logits"]

Splitting the data properly

Before anything else. Random splitting here would put photos of the same physical bottle into both train and test, and the model would score well by recognising that bottle rather than by recognising plastic.

The data pipeline

import torch, torchvision
from torch import nn
from torchvision import transforms, datasets
from torch.utils.data import DataLoader

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

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

eval_tf = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

train_ds = datasets.ImageFolder("data/train", train_tf)
val_ds   = datasets.ImageFolder("data/val",   eval_tf)
test_ds  = datasets.ImageFolder("data/test",  eval_tf)

train_dl = DataLoader(train_ds, batch_size=32, shuffle=True,  num_workers=4)
val_dl   = DataLoader(val_ds,   batch_size=64, shuffle=False, num_workers=4)

Those normalization numbers are not arbitrary. They are the channel means and standard deviations of ImageNet, and the pretrained weights were learned with inputs scaled that way. Feed raw 0–1 pixels instead and every frozen filter sees inputs outside the range it expects.

Trace one pixel. A red channel value of 200 becomes $200/255 = 0.784$ after ToTensor, then:

$\frac{0.784 - 0.485}{0.229} = \frac{0.299}{0.229} = 1.306$

Augmentation goes on train only. Validation and test use a fixed resize and centre crop, because you need the same images evaluated the same way every epoch — otherwise a metric change could be augmentation randomness rather than model improvement.

Handling class imbalance

Metal has 140 images against cardboard’s 520. Left alone, a model that predicts cardboard for everything already gets 43% accuracy, and gradient descent finds that shortcut early.

Class weights make each class contribute equally to the loss:

$w_c = \frac{N}{k \cdot n_c}$

where $N$ is the total, $k$ is the number of classes, and $n_c$ is the count for class $c$. With $N = 1200$ and $k = 4$:

The ratio between the extremes is $2.143 / 0.577 = 3.71$, which is exactly $520/140$. That is the point: the weight cancels the count.

counts = torch.tensor([520., 340., 200., 140.])
weights = counts.sum() / (len(counts) * counts)
criterion = nn.CrossEntropyLoss(weight=weights.to(device), label_smoothing=0.05)

Building the model

from torchvision.models import resnet18, ResNet18_Weights

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

for p in model.parameters():          # freeze everything
    p.requires_grad = False

for p in model.layer4.parameters():   # unfreeze the last block
    p.requires_grad = True

model.fc = nn.Linear(model.fc.in_features, 4)   # new head, trainable by default
model = model.to(device)

n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
n_total = sum(p.numel() for p in model.parameters())
print(f"training {n_train:,} of {n_total:,} parameters "
      f"({100*n_train/n_total:.1f}%)")

Replacing model.fc creates a fresh layer with requires_grad=True, so it trains even though everything else was frozen first. Order matters — freeze, then replace.

Two learning rates

The new head starts from random weights and must move a long way. The pretrained layer4 is already close to useful and should only be nudged. Give them different learning rates:

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

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=25)

The training loop

def run_epoch(loader, train):
    model.train() if train else model.eval()
    total_loss, correct, n = 0.0, 0, 0
    with torch.set_grad_enabled(train):
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            out = model(x)
            loss = criterion(out, y)
            if train:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * y.size(0)
            correct    += (out.argmax(1) == y).sum().item()
            n          += y.size(0)
    return total_loss / n, correct / n

best_val, patience, bad = 0.0, 5, 0
for epoch in range(25):
    tr_loss, tr_acc = run_epoch(train_dl, True)
    va_loss, va_acc = run_epoch(val_dl,  False)
    scheduler.step()
    print(f"epoch {epoch:02d}  train {tr_loss:.3f}/{tr_acc:.3f}  "
          f"val {va_loss:.3f}/{va_acc:.3f}")

    if va_acc > best_val:
        best_val, bad = va_acc, 0
        torch.save(model.state_dict(), "best.pt")
    else:
        bad += 1
        if bad >= patience:
            print(f"early stop at epoch {epoch}")
            break

Reading the loss curves

Switch between the scenarios in the control above. Each has a distinct signature and a different fix.

Here is what a real run on this dataset looks like:

The two curves separate at epoch 10. Before that, both fall together and the model is genuinely learning. After that the training loss keeps dropping while validation rises — the model is memorising the 840 training images. The gap between the curves is the amount of memorisation.

The confusion matrix, and what to do with it

model.load_state_dict(torch.load("best.pt"))
model.eval()

from sklearn.metrics import confusion_matrix, classification_report
preds, trues = [], []
with torch.no_grad():
    for x, y in test_dl:
        preds += model(x.to(device)).argmax(1).cpu().tolist()
        trues += y.tolist()

print(confusion_matrix(trues, preds))
print(classification_report(trues, preds, target_names=train_ds.classes))

Overall accuracy is 80.6%, which tells you nothing you can act on. The matrix does. Nineteen of the 35 errors are the glass–plastic pair: 11 glass called plastic, 8 plastic called glass. That single pair is 54% of all mistakes.

Now look at the images. Clear plastic bottles and clear glass bottles look nearly identical in a photo — the real distinguishing cues are weight and the sound they make, neither of which is in the image. What is in the image: glass has sharper specular highlights and slightly different edge refraction.

This is the workflow from error analysis applied to a real model. The instinct is to reach for a bigger architecture; the matrix says collect glass photos instead, for four times the gain.

Practice task

Use any four-class image dataset, or photograph 300 items around your home.

1. Split by session or object, never at random. Write down the rule you used.
2. Train with the whole backbone frozen, head only. Record test accuracy.
3. Unfreeze layer4 at 1e-4 and retrain. Record the change.
4. Train from scratch with weights=None for the same number of epochs. Record it.
5. Plot train and validation loss for all three runs on one chart.
6. Print the confusion matrix for your best model and identify the dominant error pair.
7. Look at 20 images from that pair. Write one sentence on why the model confuses them.
8. Now redo step 1 with a random split and retrain. Compare test accuracy.

Step 8 is deliberately the last one. The random-split model will score several points higher on a test set it has effectively already seen. Watching that inflation happen on your own data is the most reliable way to never make the mistake again.

Summary

Transfer learning works because early network layers learn edges and textures that are the same for every photograph. Training only a new head means fitting $512 \times 4 + 4 = 2052$ parameters instead of 11.7 million, which is a sane ratio for 1,200 images. Unfreezing the last block adds a few points; training from scratch on this data loses about 40.

Split by session so no physical object appears twice. Normalise with ImageNet statistics, so a red value of 200 becomes 1.306. Augment train only. Weight the loss by $N/(k \cdot n_c)$, which gave 0.577 for cardboard and 2.143 for metal — a 3.71× ratio exactly cancelling the 520:140 imbalance.

Watch the two loss curves. They separated at epoch 10 in the run above, and everything after that was memorisation. Save on every validation improvement so you keep the best model rather than the last one.

Finally, read the confusion matrix rather than the accuracy. 80.6% overall was not actionable; “19 of 35 errors are glass versus plastic” was, and it pointed at collecting glass photos rather than at a bigger architecture.

What comes next

Overfitting appeared at epoch 10, and the standard first response is better data handling. Vision data pipelines and augmentation covers which augmentations actually help, which quietly destroy your labels, and how to build an input pipeline that does not starve the GPU. After that, CNN architectures from LeNet to ResNet explains what is inside the backbone you have been using as a black box.