Project: Building an Image Classifier

Build the mental model

A working image classifier isn't one clever trick — it's a full pipeline where every stage depends on the one before it holding its shape correctly. This project wires the whole thing together: a Dataset class that generates labeled synthetic images (bright-top vs bright-bottom patterns standing in for real class differences), a DataLoader that batches and shuffles them, a small CNN with two Conv2d+ReLU+MaxPool blocks that progressively shrink spatial size while growing channel depth, and a Linear head that turns the final feature map into class scores. Training loops over this with CrossEntropyLoss and Adam, and evaluation runs the same forward pass with gradients disabled on held-out data. The naive alternative — flattening pixels straight into a stack of Linear layers — throws away the 2D spatial structure that makes images images: a CNN's convolutions detect local patterns (edges, blobs) regardless of where they sit in the frame, and pooling keeps that recognition even if the pattern shifts slightly. Swap the synthetic generator for CIFAR-10 or ImageNet loading and nothing else in this pipeline changes — that's the whole point of building it this way.

Connect it to a real scenario

This exact pipeline is what would power a 'suggested lesson thumbnail' or 'auto-tag diagram type' feature on the Tutorial Platform — feeding screenshots or diagram images from lesson content through a CNN to classify them (e.g., 'code screenshot' vs 'architecture diagram' vs 'chart') for automatic tagging and search filtering. The Dataset/DataLoader/train/evaluate structure here is identical to what a real feature would use; only the image source changes from synthetic tensors to actual uploaded lesson images, and the classes change from two toy patterns to real content categories the platform needs to distinguish.

A Dataset + DataLoader pair standardizes how any project feeds batched data to a model, synthetic or real.
Conv2d+ReLU+MaxPool blocks build spatial pattern detectors that generalize better than flattening pixels into Linear layers.
Falling training loss plus a separate held-out accuracy check together prove the model generalizes, not just memorizes.

Try the working example

python

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

class SyntheticImageDataset(Dataset):
    def __init__(self, n_samples=200, size=16):
        self.data, self.labels = [], []
        for i in range(n_samples):
            label = i % 2
            img = torch.rand(1, size, size) * 0.3
            if label == 1:
                img[:, :size // 2, :] += 0.6
            else:
                img[:, size // 2:, :] += 0.6
            self.data.append(img)
            self.labels.append(label)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc = nn.Linear(16 * 4 * 4, 2)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))  # 16 -> 8
        x = self.pool(torch.relu(self.conv2(x)))  # 8 -> 4
        x = x.view(x.size(0), -1)
        return self.fc(x)

train_loader = DataLoader(SyntheticImageDataset(200), batch_size=16, shuffle=True)
test_loader = DataLoader(SyntheticImageDataset(40), batch_size=16)

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

for epoch in range(5):
    total_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss / len(train_loader):.4f}")

model.eval()
correct, total = 0, 0
with torch.no_grad():
    for images, labels in test_loader:
        preds = model(images).argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
print(f"Test accuracy: {correct/total:.2%}")

You should see

The per-epoch loss steadily decreases across the 5 epochs, and the script finishes by printing a Test accuracy percentage (typically well above 90%) on the held-out synthetic test set.

5-minute try-it

Change conv1's output channels from 8 to 32 and increase training to 10 epochs — how does accuracy change? Then create a new synthetic pattern that splits bright/dark left-right instead of top-bottom, and check whether the same model architecture learns it just as easily.

One important caution

Forgetting model.eval() (and torch.no_grad()) during evaluation lets dropout/batchnorm behave as if still training and wastes memory tracking gradients you'll never use.

Mismatching the flattened feature size passed to the first Linear layer after conv/pool layers — miscounting the spatial dimensions after pooling is one of the most common CNN shape-mismatch bugs.

PyTorch Examples — MNIST — Deep Learning