Apply CNNs for Image Classification

You've built pipelines for structured outputs and retrieval, but now you're staring at a folder of images that needs categorization—manually labeling thousands of photos is a nightmare, and traditional code can't 'see' the patterns. This is the exact pain point that convolutional neural networks (CNNs) solve: they automate feature extraction from raw pixels, letting you classify images with human-level accuracy. In this lesson, you'll go from zero to a working CNN classifier using PyTorch, learning the core mechanics and how to apply them hands-on.

The problem this lesson solves

Classifying images with classical machine learning is brutally inefficient. If you try to feed raw pixel values into a logistic regression or a random forest, you'll hit two walls:

• Curse of dimensionality — a tiny 64×64 color image has 12,288 features per sample. With thousands of images, training becomes slow and prone to overfitting.
• No spatial awareness — flattening pixels destroys the 2D structure (edges, textures, shapes). The model can't learn that a cat's ear is near its head.

A CNN solves this by learning hierarchical features automatically: early layers detect edges and colors, middle layers assemble them into shapes, and final layers recognize complete objects like 'cat' or 'dog'. You no longer hand-engineer features like HOG or SIFT—the network learns them from data.

Why now? With libraries like PyTorch and TensorFlow, building a CNN takes fewer than 50 lines of code. The barrier is no longer math—it's knowing how to structure data and train effectively.

Core concept / mental model

Think of a CNN as a stack of magnifying glasses. Each layer looks at a slightly larger region of the image and builds a more abstract representation:

• Convolutional layers slide small filters (like 3×3 squares) across the image to detect patterns. Imagine scanning a photo with a magnifying glass that highlights vertical edges.
• Activation functions (like ReLU) add non-linearity, letting the network learn complex patterns.
• Pooling layers shrink the image (e.g., take the maximum value in a 2×2 block), keeping important features while reducing computation.

By the time the data reaches the fully connected layer at the end, it's a compact vector of learned features. A softmax layer then outputs probabilities for each class.

Key terms: - Filter/kernel — a small matrix that slides over the image. - Stride — how many pixels the filter moves each step. - Padding — adding border zeros to preserve spatial dimensions. - Feature map — the output of a convolution operation.

How it works step by step

Applying a CNN to image classification follows a repeatable pipeline. Here's the logical order:

1. Load data — download and organize images into train/validation/test sets with labels.
2. Preprocess images — resize to a uniform size, convert to tensors, normalize pixel values (e.g., mean 0.5, std 0.5).
3. Define the CNN architecture — choose number of conv layers, filters, pooling, and fully connected layers.
4. Set loss function and optimizer — cross-entropy loss for multi-class classification, Adam or SGD optimizer.
5. Train the model — iterate over mini-batches, compute gradients, update weights.
6. Evaluate — compute accuracy and loss on validation and test sets.

Each step feeds into the next. If your accuracy is low, you tweak the architecture or training hyperparameters.

Hands-on walkthrough

Let's implement a CNN for the classic CIFAR-10 dataset (60,000 32×32 color images across 10 classes). You'll need PyTorch installed (pip install torch torchvision).

Step 1: Load and preprocess data

import torch
import torchvision
import torchvision.transforms as transforms

# Define transforms: convert to tensor and normalize (mean, std)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# Download and load training and test sets
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64,
                                         shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
print(f"Training samples: {len(trainset)}, Test samples: {len(testset)}")

Expected output:

Training samples: 50000, Test samples: 10000

Step 2: Define the CNN architecture

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

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)  # input: 3 channels (RGB), output: 16 feature maps
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)  # halves spatial size
        self.fc1 = nn.Linear(32 * 8 * 8, 256)  # after 2 pools: 32x32 -> 8x8
        self.fc2 = nn.Linear(256, 10)  # 10 classes

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 32 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleCNN()
print(model)

Expected output: summary of layers (conv1, conv2, pool, fc1, fc2).

Step 3: Train the model

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(5):  # 5 epochs for demo
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(trainloader):.4f}")

print("Finished training")

Expected output: loss values decreasing each epoch, e.g., Epoch 1, Loss: 1.89 then Epoch 5, Loss: 1.12 (exact numbers vary).

Step 4: Evaluate accuracy

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Test accuracy: {100 * correct / total:.2f}%")

Expected output: something like Test accuracy: 62.35% — not world-class, but a working baseline.

Pro tip: For production, you'd train for 20+ epochs, use data augmentation (random flips, crops), and consider a pretrained ResNet.

Compare options / when to choose what

You might wonder: should I use a custom CNN, a pretrained model, or a fully connected network? Here's a quick comparison:

Variations to know: - Transfer learning — take a pretrained model, freeze early layers, retrain last layers on your data. Hugely effective. - Data augmentation — flip, rotate, or crop images during training to improve generalization. - Different pooling — global average pooling instead of max pooling is common in modern CNNs.

For most real-world tasks, pretrained models with transfer learning win. Build a custom CNN only when you need extreme customization or studying internals.

Troubleshooting & edge cases

Even with the right code, things go wrong. Here are common issues and fixes:

• Loss not decreasing — likely a too-high learning rate or a bug in model definition (e.g., wrong input size). Lower LR, check tensor shapes.
• Out-of-memory (OOM) — reduce batch size or image resolution.
• Validation accuracy much lower than training — overfitting. Add dropout, regularization, or data augmentation.
• Wrong input size error in forward pass — recalculate the flattened size after convolutions/pooling. Break after conv layers and print x.shape.
• CUDA errors — ensure tensors and model are on the same device; use .to(device).

Edge case: Grayscale images (1 channel) need transforms.Grayscale() and a Conv2d with in_channels=1.

What you learned & what's next

You now know how to apply CNNs for image classification: you understand the mental model of convolutional layers, pooling, and fully connected outputs; you implemented a complete pipeline in PyTorch from loading CIFAR-10 to evaluating accuracy; and you can troubleshoot common failures like shape mismatches and overfitting. You've achieved the learning objective of explaining and practicing CNN classification.

Next up in the trAI path: You'll move from classifying static images to more advanced tasks like object detection or image segmentation, where CNNs form the backbone. Or you may explore hyperparameter tuning to push accuracy above 80%.

Keep the momentum — take your trained model and try to classify a single image from the internet, then share your results!