A complete image classification project — using transfer learning to build a solid classifier without needing a massive dataset or huge compute budget, the standard first real computer vision project.
Problem Statement
Build an image classifier that distinguishes between 10 categories of everyday objects, achieving at least 85% test accuracy, with a training pipeline that could realistically be adapted to a different image classification task by swapping the dataset.
Dataset
CIFAR-10 (60,000 32×32 color images across 10 classes) is a good starting dataset — small enough to train quickly, standard enough that results are easy to sanity-check against known benchmarks, and available directly through torchvision.datasets.
Architecture & Approach
Rather than training a CNN from scratch (which would need far more data and compute to reach strong accuracy), this project uses transfer learning — starting from a ResNet pretrained on ImageNet, replacing its final classification layer, and fine-tuning on CIFAR-10, directly applying the Transfer Learning category's concepts.
Step-by-Step Build
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as T
# 1. Data pipeline -- resize to match ImageNet-pretrained expectations, normalize accordingly
transform = T.Compose([
T.Resize(224),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True, num_workers=4)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=64, shuffle=False, num_workers=4)
# 2. Model -- pretrained ResNet18 with a replaced final layer
model = torchvision.models.resnet18(weights='IMAGENET1K_V1')
for param in model.parameters():
param.requires_grad = False # freeze the pretrained backbone initially
model.fc = nn.Linear(model.fc.in_features, 10) # new head, trainable by default
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 3. Train only the new head first
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(5):
model.train()
for x_batch, y_batch in train_loader:
x_batch, y_batch = x_batch.to(device), y_batch.to(device)
optimizer.zero_grad()
loss = loss_fn(model(x_batch), y_batch)
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1} done, last batch loss: {loss.item():.4f}")
# 4. Evaluate
model.eval()
correct, total = 0, 0
with torch.no_grad():
for x_batch, y_batch in test_loader:
x_batch, y_batch = x_batch.to(device), y_batch.to(device)
preds = model(x_batch).argmax(dim=1)
correct += (preds == y_batch).sum().item()
total += y_batch.size(0)
print(f"Test accuracy: {correct/total:.4f}")Expected Results
Training only the new classification head (backbone frozen) for a few epochs should already reach roughly 80-85% test accuracy on CIFAR-10 — unfreezing and fine-tuning the last few backbone layers with a smaller learning rate for additional epochs typically pushes this further, illustrating the two-stage fine-tuning approach from Partial vs Full Fine-Tuning.
Key Learnings & Extensions
- This project directly demonstrates why transfer learning is so valuable for smaller datasets — training a comparable CNN from scratch on CIFAR-10 alone typically needs many more epochs and careful tuning to reach similar accuracy.
- Extension: Add data augmentation (random crop, flip) and compare accuracy with and without it — a direct, hands-on illustration of Data Augmentation Pipeline's value.
- Extension: Perform error analysis — look at the confusion matrix and inspect specific misclassified images, following the process in Error Analysis.
