Image Classification with CNNs — Deep Learning Project

What You'll Build

A complete image classification workflow in PyTorch that:

Loads CIFAR-10 (60,000 32×32 colour images, 10 classes) with proper data augmentation
Trains a custom 3-layer CNN from scratch (~78% accuracy)
Fine-tunes a pre-trained ResNet-18 (~92% accuracy)
Plots training curves and a per-class confusion matrix
Deploys both models in a Gradio web app

Project Setup

python -m venv .venv
source .venv/bin/activate

pip install torch torchvision matplotlib seaborn gradio tqdm

GPU check: Run torch.cuda.is_available() — if True, training will be much faster. Google Colab provides a free T4 GPU: Runtime → Change runtime type → T4 GPU.

Step 1 — Load and Visualize CIFAR-10

import torch
import torchvision
import torchvision.transforms as T
import matplotlib.pyplot as plt
import numpy as np

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
CLASSES = ('plane','car','bird','cat','deer',
           'dog','frog','horse','ship','truck')

# Training transforms with augmentation
train_transforms = T.Compose([
    T.RandomHorizontalFlip(),
    T.RandomCrop(32, padding=4),
    T.ColorJitter(brightness=0.2, contrast=0.2),
    T.ToTensor(),
    T.Normalize((0.4914, 0.4822, 0.4465),
                (0.2470, 0.2435, 0.2616)),
])
val_transforms = T.Compose([
    T.ToTensor(),
    T.Normalize((0.4914, 0.4822, 0.4465),
                (0.2470, 0.2435, 0.2616)),
])

trainset = torchvision.datasets.CIFAR10('./data', train=True,
                                         download=True, transform=train_transforms)
valset   = torchvision.datasets.CIFAR10('./data', train=False,
                                         download=True, transform=val_transforms)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                           shuffle=True, num_workers=2)
valloader   = torch.utils.data.DataLoader(valset, batch_size=256,
                                           shuffle=False, num_workers=2)
print(f"Train: {len(trainset):,}  Val: {len(valset):,}")

# Visualize a batch
imgs, labels = next(iter(trainloader))
imgs = imgs[:16]
grid = torchvision.utils.make_grid(imgs, nrow=8, normalize=True)
plt.figure(figsize=(14, 4))
plt.imshow(grid.permute(1,2,0))
plt.title([CLASSES[l] for l in labels[:16]])
plt.axis('off')
plt.show()

Step 2 — Custom CNN from Scratch

import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.bn1   = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2   = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn3   = nn.BatchNorm2d(128)
        self.pool  = nn.MaxPool2d(2)
        self.drop  = nn.Dropout(0.4)
        self.fc1   = nn.Linear(128 * 4 * 4, 256)
        self.fc2   = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))  # 32→16
        x = self.pool(F.relu(self.bn2(self.conv2(x))))  # 16→8
        x = self.pool(F.relu(self.bn3(self.conv3(x))))  # 8→4
        x = x.view(x.size(0), -1)
        x = self.drop(F.relu(self.fc1(x)))
        return self.fc2(x)

cnn = SimpleCNN().to(DEVICE)
print(f"Parameters: {sum(p.numel() for p in cnn.parameters()):,}")

Step 3 — Training Loop

import torch.optim as optim
from tqdm import tqdm

def train_model(model, trainloader, valloader, epochs=20, lr=1e-3):
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    criterion = nn.CrossEntropyLoss()

    history = {'train_loss': [], 'val_acc': []}

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for imgs, labels in tqdm(trainloader, desc=f'Epoch {epoch+1}/{epochs}', leave=False):
            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
            optimizer.zero_grad()
            loss = criterion(model(imgs), labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        # Validation
        model.eval()
        correct = total = 0
        with torch.no_grad():
            for imgs, labels in valloader:
                imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
                preds = model(imgs).argmax(dim=1)
                correct += (preds == labels).sum().item()
                total   += labels.size(0)
        acc = correct / total

        history['train_loss'].append(total_loss / len(trainloader))
        history['val_acc'].append(acc)
        scheduler.step()
        print(f'Epoch {epoch+1:2d}  Loss: {history["train_loss"][-1]:.3f}  Val Acc: {acc:.4f}')

    return history

cnn_history = train_model(cnn, trainloader, valloader, epochs=20)

After 20 epochs you should see ~75–80% validation accuracy.

Step 4 — Transfer Learning with ResNet-18

import torchvision.models as models

resnet = models.resnet18(weights='IMAGENET1K_V1')

# Freeze all layers except the final classifier
for param in resnet.parameters():
    param.requires_grad = False

# Replace the head for 10-class CIFAR-10
resnet.fc = nn.Linear(resnet.fc.in_features, 10)
resnet = resnet.to(DEVICE)

# ResNet expects 224×224 — use a resize transform for fine-tuning
train_t_resnet = T.Compose([
    T.Resize(64),   # upscale from 32 to 64 (lighter than 224)
    T.RandomHorizontalFlip(),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_t_resnet = T.Compose([
    T.Resize(64),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
trainset.transform = train_t_resnet
valset.transform   = val_t_resnet

# Train only the new head first (5 epochs)
resnet_history = train_model(resnet, trainloader, valloader, epochs=5, lr=1e-3)

# Then unfreeze all layers and fine-tune with smaller LR (10 more epochs)
for param in resnet.parameters():
    param.requires_grad = True
resnet_history2 = train_model(resnet, trainloader, valloader, epochs=10, lr=1e-4)

Two-phase fine-tuning: Train only the head first so the randomly initialized weights don't corrupt the pre-trained features. Once the head converges, unfreeze everything and fine-tune end-to-end at a much lower learning rate.

Step 5 — Compare Training Curves

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

axes[0].plot(cnn_history['val_acc'], label='CNN from scratch')
axes[0].plot(resnet_history['val_acc'] + resnet_history2['val_acc'], label='ResNet-18')
axes[0].set_title('Validation Accuracy'); axes[0].set_xlabel('Epoch')
axes[0].legend(); axes[0].grid(True, alpha=0.3)

axes[1].plot(cnn_history['train_loss'], label='CNN from scratch')
axes[1].set_title('Training Loss'); axes[1].set_xlabel('Epoch')
axes[1].legend(); axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

Step 6 — Per-Class Analysis

from sklearn.metrics import ConfusionMatrixDisplay
import numpy as np

all_preds, all_labels = [], []
resnet.eval()
with torch.no_grad():
    for imgs, labels in valloader:
        imgs = imgs.to(DEVICE)
        preds = resnet(imgs).argmax(dim=1).cpu()
        all_preds.extend(preds.numpy())
        all_labels.extend(labels.numpy())

ConfusionMatrixDisplay.from_predictions(
    all_labels, all_preds,
    display_labels=CLASSES,
    cmap='Blues', xticks_rotation=45
)
plt.title('ResNet-18 Confusion Matrix (CIFAR-10)')
plt.tight_layout()
plt.show()

Common mistakes to look for: cats ↔ dogs, cars ↔ trucks, deer ↔ horse. These are visually similar and consistently confuse even SOTA models at low resolution.

Step 7 — Gradio Web App

import gradio as gr
from PIL import Image

# Save the models
torch.save(cnn.state_dict(),    'cnn_cifar10.pt')
torch.save(resnet.state_dict(), 'resnet_cifar10.pt')

preprocess = val_t_resnet

def predict(image, model_choice):
    model = resnet if model_choice == "ResNet-18" else cnn
    img_t = preprocess(image).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        logits = model(img_t)
        probs  = torch.softmax(logits, dim=1)[0]
    return {CLASSES[i]: float(probs[i]) for i in range(10)}

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload an image"),
        gr.Radio(["ResNet-18", "Custom CNN"], value="ResNet-18", label="Model"),
    ],
    outputs=gr.Label(num_top_classes=5, label="Predictions"),
    title="CIFAR-10 Image Classifier",
    description="Upload any image — the model predicts which of the 10 CIFAR-10 classes it belongs to.",
)

demo.launch()

Results Summary

Model	Val Accuracy	Parameters	Training Time (GPU)
Custom CNN (from scratch)	~78%	~1.2M	~4 min
ResNet-18 (fine-tuned)	~92%	11.2M	~8 min

What to Try Next

EfficientNet: Swap ResNet-18 for EfficientNet-B0 — similar accuracy with fewer parameters
Mixup / CutMix: Advanced data augmentation that blends two images for better generalization
Learning rate finder: Use torch-lr-finder to find the optimal LR automatically
CIFAR-100: Scale up to 100 classes — harder, and forces you to deal with class imbalance and longer training