Mini Project: Image Classification with PyTorch

Site Admin · 11 Sep 2026 · 7 views

Mini Project: Image Classification with PyTorch

This project classifies handwritten digits from MNIST, the classic entry task for computer vision: 60,000 training and 10,000 test images of 28 by 28 grayscale digits. Assembled here are every tool from this tutorial in one script. Reuse the SmallCNN class from the CNN post, or paste the minimal definition into the same cell.

Load the Data

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

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
train = datasets.MNIST(root="data", train=True, download=True, transform=transform)
test = datasets.MNIST(root="data", train=False, download=True, transform=transform)
loader = DataLoader(train, batch_size=64, shuffle=True)

ToTensor converts pixel values to tensors in the zero-to-one range; Normalize rescales around zero, which helps every optimizer.

Build and Train

import torch.nn as nn
import torch.optim as optim

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

for epoch in range(3):
    for images, labels in loader:
        optimizer.zero_grad()
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()
    print("epoch", epoch + 1, "loss", round(loss.item(), 4))

Evaluate on Unseen Digits

from sklearn.metrics import accuracy_score

preds, y_true = [], []
for images, labels in test_loader:
    logits = model(images)
    preds.extend(logits.argmax(dim=1).tolist())
    y_true.extend(labels.tolist())
print("Test accuracy:", round(accuracy_score(y_true, preds), 4))

Expect a test accuracy above 90 percent after only three epochs even with the small CNN, with tighter wins following from a deeper model and more epochs.