Getting Started with Computer Vision Using PyTorch

Computer vision is one of the most exciting areas in machine learning. Being able to write code that "sees" — classifying images, detecting objects, reading handwriting — feels like actual magic the first time it works. And with PyTorch, the gap between understanding the concepts and running working code is smaller than you might think. This post gets you from zero to a trained image classifier.

What is Computer Vision?

Computer Vision (CV) is the field that enables machines to interpret and understand visual information from the world. Core tasks include:

Image Classification — "What is in this image?" (our focus today)
Object Detection — "Where are the objects, and what are they?"
Semantic Segmentation — "Which pixels belong to which object?"
Image Generation — GANs, diffusion models, etc.

We'll focus on image classification — it's the foundation everything else builds on.

The MNIST Dataset

MNIST is the "hello world" of computer vision: 70,000 handwritten digit images (0–9), each 28×28 pixels in greyscale. It's small enough to train on a CPU in minutes but real enough to teach core concepts.

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

# ── Transforms: normalize pixel values to mean=0.5, std=0.5 ──────
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# ── Download & load ───────────────────────────────────────────────
train_dataset = datasets.MNIST(root='./data', train=True,
                                download=True, transform=transform)
val_dataset   = datasets.MNIST(root='./data', train=False,
                                download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader   = DataLoader(val_dataset,   batch_size=64)

print(f"Train samples: {len(train_dataset)}")  # 60000
print(f"Val samples:   {len(val_dataset)}")    # 10000

Why CNNs? The Problem With Fully-Connected Networks

A 28×28 image has 784 pixels. If you flatten it and feed it to an MLP, each neuron in the first layer needs a weight for every pixel — 784 weights per neuron. For a 256-neuron first layer that's already 200,000+ parameters, and the network has no idea that adjacent pixels are spatially related.

Convolutional Neural Networks (CNNs) solve both problems:

Weight sharing: the same small filter (e.g. 3×3 kernel) slides across the whole image — far fewer parameters
Local connectivity: each neuron only sees a small patch of the image, learning local features (edges, textures) that compose into higher-level patterns in deeper layers

Building a CNN in PyTorch

class DigitCNN(nn.Module):
    def __init__(self):
        super().__init__()
        # Block 1: 1 input channel (greyscale) → 32 feature maps
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),  # → (32, 28, 28)
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),                             # → (32, 14, 14)
        )
        # Block 2: 32 → 64 feature maps
        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, padding=1), # → (64, 14, 14)
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),                             # → (64, 7, 7)
        )
        # Classifier head
        self.classifier = nn.Sequential(
            nn.Flatten(),               # → (64 * 7 * 7 = 3136,)
            nn.Linear(3136, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 10),         # 10 digit classes
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return self.classifier(x)

model = DigitCNN()
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
# → ~220,000 — far fewer than an equivalent MLP

Key Components Explained

Conv2d(in, out, kernel_size=3, padding=1) — applies out different 3×3 filters; padding=1 keeps spatial dimensions the same
BatchNorm2d — normalises feature maps across the batch; speeds training and acts as mild regularisation
MaxPool2d(2) — halves spatial dimensions by taking the max value in each 2×2 window (downsampling + translation invariance)
Dropout(0.3) — randomly zeroes 30 % of activations during training to prevent overfitting

Training the CNN

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model  = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

def evaluate(loader):
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            correct += (outputs.argmax(1) == labels).sum().item()
            total   += labels.size(0)
    return correct / total

for epoch in range(1, 11):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        loss = criterion(model(images), labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    scheduler.step()

    val_acc = evaluate(val_loader)
    print(f"Epoch {epoch:2d}  loss={running_loss/len(train_loader):.4f}  val_acc={val_acc:.4f}")

# Epoch 10  loss=0.0312  val_acc=0.9921  ← 99.2% on MNIST!

Transfer Learning: Standing on Giants' Shoulders

Training from scratch works for simple datasets like MNIST. For complex real-world images (medical scans, product photos, satellite imagery), you rarely have enough labelled data. Transfer learning is the solution: start from a model pre-trained on ImageNet (1.2 million images, 1000 classes) and fine-tune it on your data.

from torchvision import models

# Load a pre-trained ResNet-18
backbone = models.resnet18(weights='IMAGENET1K_V1')

# Freeze all layers — we only want to fine-tune the classifier head
for param in backbone.parameters():
    param.requires_grad = False

# Replace the final classification layer for our task (e.g., 5 classes)
num_classes = 5
backbone.fc = nn.Linear(backbone.fc.in_features, num_classes)

# Now only backbone.fc parameters will be updated during training
optimizer = optim.Adam(backbone.fc.parameters(), lr=1e-3)

This approach often reaches 90 %+ accuracy on small datasets (a few hundred images per class) that would be impossible to learn from scratch.

Data Augmentation — Free Performance

Augmentation artificially expands your training set by applying random transformations that don't change the label:

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(degrees=10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomCrop(224, padding=10),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),   # ImageNet stats
])

A model trained with augmentation generalises much better than one trained on the raw images.

What to Try Next

CIFAR-10 — 10 classes of 32×32 colour images; harder than MNIST, still manageable
Your own dataset — collect ~200 images per class, set up train/val folders, apply transfer learning
Object detection — explore YOLO or Faster R-CNN for localising objects
Grad-CAM — visualise which pixels your model is "looking at" when it makes a prediction

Conclusion

Computer vision went from a niche research topic to a practical engineering skill. With PyTorch, torchvision, and a bit of patience, you can train image classifiers, apply transfer learning from powerful backbones, and build CV systems that work on real data. Start with MNIST, graduate to CIFAR-10, then pick a problem that genuinely interests you.

For the theory behind the neural network building blocks used here, see Understanding Neural Networks: A Practical Introduction.