Chapter 5: Practical Project

1. Implementing an Image Classification Task

Image Classification is the task of predicting which class (category) an input image belongs to. In this chapter, we'll build a model that classifies grayscale images containing three types of patterns (horizontal stripes, vertical stripes, and a diagonal cross) into three classes. The overall workflow is as follows.

1.1 Generating Synthetic Image Data

First, let's write a function that uses numpy to generate synthetic images for 3 classes. Each class embeds a faint pattern (horizontal stripes, vertical stripes, or a diagonal cross) within noise.

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision import transforms

torch.manual_seed(42)
np.random.seed(42)

def generate_synthetic_images(num_samples_per_class=150, image_size=28,
                               num_classes=3, seed=42):
    """Generate synthetic grayscale images for 3 classes.

    label=0: horizontal stripes, label=1: vertical stripes, label=2: diagonal cross
    """
    rng = np.random.default_rng(seed)
    images = []
    labels = []
    for label in range(num_classes):
        for _ in range(num_samples_per_class):
            img = rng.normal(loc=0.15, scale=0.20,
                              size=(image_size, image_size)).astype(np.float32)
            if label == 0:  # horizontal stripes
                img[8:12, :] += 0.35
                img[16:20, :] += 0.35
            elif label == 1:  # vertical stripes
                img[:, 8:12] += 0.35
                img[:, 16:20] += 0.35
            else:  # diagonal cross
                idx = np.arange(image_size)
                img[idx, idx] += 0.35
                img[idx, image_size - 1 - idx] += 0.35
            img = np.clip(img, 0.0, 1.0)
            images.append(img)
            labels.append(label)

    images = np.stack(images).astype(np.float32)[:, np.newaxis, :, :]  # (N, 1, H, W)
    labels = np.array(labels, dtype=np.int64)

    # shuffle
    perm = rng.permutation(len(labels))
    return torch.from_numpy(images[perm]), torch.from_numpy(labels[perm])

all_images, all_labels = generate_synthetic_images()
print(f"images: {all_images.shape}, labels: {all_labels.shape}")

Example output:

images: torch.Size([450, 1, 28, 28]), labels: torch.Size([450])

Each image has the shape (channels, height, width) = (1, 28, 28), a grayscale image. This is the same format used by real image datasets such as MNIST, so the code that follows can be applied to real data as-is.

1.2 Creating a Custom Dataset and DataLoader

So we can apply data augmentation (explained in the next section) on a per-sample basis, let's define a custom class that inherits from torch.utils.data.Dataset.

class SyntheticImageDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform

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

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

# Split into training, validation, and test sets (70% / 15% / 15%)
base_dataset = SyntheticImageDataset(all_images, all_labels, transform=None)
n_total = len(base_dataset)
n_train = int(n_total * 0.7)
n_val = int(n_total * 0.15)
n_test = n_total - n_train - n_val

train_subset, val_subset, test_subset = random_split(
    base_dataset, [n_train, n_val, n_test],
    generator=torch.Generator().manual_seed(42)
)
print(f"train: {n_train}, val: {n_val}, test: {n_test}")

Example output:

train: 315, val: 67, test: 68

Because random_split splits the Dataset by index, the samples in the training, validation, and test subsets are guaranteed not to overlap. After splitting, let's create DataLoaders so we can retrieve data in batches.

train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_subset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_subset, batch_size=32, shuffle=False)

images_batch, labels_batch = next(iter(train_loader))
print(f"batch images: {images_batch.shape}, batch labels: {labels_batch.shape}")

Example output:

batch images: torch.Size([32, 1, 28, 28]), batch labels: torch.Size([32])

2. Building a CNN Model

A Convolutional Neural Network (CNN) is a neural network architecture that excels at handling grid-like data such as images. It's typically built by combining the following three types of layers.

• Convolutional Layer (nn.Conv2d): slides a small filter (kernel) across the image to detect local features such as edges and patterns
• Pooling Layer (nn.MaxPool2d): reduces the spatial size of feature maps, lowering computational cost while increasing robustness to positional shifts
• Fully Connected Layer (nn.Linear): flattens the features extracted by the convolutions and performs the final class prediction

Let's start by defining a simple CNN.

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=3):
        super().__init__()
        # Convolutional layers: 1 channel (grayscale) -> 16 channels -> 32 channels
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.relu = nn.ReLU()

        # 28x28 -> pool -> 14x14 -> pool -> 7x7, so 32*7*7 is the input dimension to the fully connected layer
        self.fc1 = nn.Linear(32 * 7 * 7, 64)
        self.fc2 = nn.Linear(64, num_classes)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))  # (B, 16, 14, 14)
        x = self.pool(self.relu(self.conv2(x)))  # (B, 32, 7, 7)
        x = x.view(x.size(0), -1)                 # flatten: (B, 32*7*7)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleCNN(num_classes=3)
sample_output = model(images_batch)
print(f"output shape: {sample_output.shape}")

Example output:

output shape: torch.Size([32, 3])

The input images are 28×28, but two rounds of pooling shrink the spatial size from 28 to 14 to 7. Tracking this shape transformation is one of the trickiest parts of CNN design. The input dimension of fc1 must always match the output shape of the final convolutional layer (channels × height × width).

3. Data Augmentation and Batch Normalization

3.1 Data Augmentation

Data Augmentation applies random transformations (flips, rotations, translations, etc.) to the training data, effectively increasing the apparent quantity and variety of data. This suppresses Overfitting — where the model memorizes the fine details of the training data — and improves generalization performance on unseen data.

Using torchvision.transforms, you can easily apply data augmentation to Tensor images.

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),   # flip horizontally with 50% probability
    transforms.RandomRotation(degrees=10),    # rotate randomly within ±10 degrees
    transforms.Normalize(mean=[0.2], std=[0.3]),  # normalize pixel values
])

# For validation/test data, apply normalization only — no randomness
eval_transform = transforms.Compose([
    transforms.Normalize(mean=[0.2], std=[0.3]),
])

3.2 Batch Normalization

Batch Normalization normalizes the output of each layer to have a mean of 0 and a variance of 1, computed per mini-batch. It stabilizes training, speeds up convergence, and enables the use of higher learning rates. For image data, it's common to insert nn.BatchNorm2d right after a convolutional layer.

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=3):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)
        self.fc1 = nn.Linear(32 * 7 * 7, 64)
        self.fc2 = nn.Linear(64, num_classes)

    def forward(self, x):
        x = self.pool(self.relu(self.bn1(self.conv1(x))))  # (B, 16, 14, 14)
        x = self.pool(self.relu(self.bn2(self.conv2(x))))  # (B, 32, 7, 7)
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)  # further suppresses overfitting in the fully connected layer
        x = self.fc2(x)
        return x

The standard pattern for ordering convolutional layers is "Conv2d → BatchNorm2d → activation function → MaxPool2d." In addition, inserting Dropout right before the fully connected layer reduces reliance on specific neurons, further improving generalization.

Finally, let's recreate the Datasets, applying train_transform (with augmentation) to the training subset and eval_transform to the validation/test subsets.

def attach_transform(subset, transform):
    """Reattach a Dataset with the specified transform to a Subset obtained from random_split."""
    subset.dataset = SyntheticImageDataset(all_images, all_labels, transform=transform)
    return subset

train_subset = attach_transform(train_subset, train_transform)
val_subset = attach_transform(val_subset, eval_transform)
test_subset = attach_transform(test_subset, eval_transform)

train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_subset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_subset, batch_size=32, shuffle=False)

4. Visualizing and Monitoring Training

To confirm that training is progressing correctly, it's essential to record and visualize the Loss and Accuracy for each epoch. First, let's define functions that perform one epoch of training and evaluation.

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss, correct, total = 0.0, 0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * images.size(0)
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
    return total_loss / total, correct / total

@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        total_loss += loss.item() * images.size(0)
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
    return total_loss / total, correct / total

Switching between model.train() and model.eval() matters a great deal. BatchNorm2d and Dropout behave differently during training and evaluation, so you must always call model.eval() before evaluating. Also, since gradient computation isn't needed during evaluation, we wrap it in @torch.no_grad() to avoid wasted computation and memory usage.

Next, let's run the full training loop and record the history in a history dictionary.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN(num_classes=3).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
num_epochs = 8

for epoch in range(num_epochs):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)

    history["train_loss"].append(train_loss)
    history["train_acc"].append(train_acc)
    history["val_loss"].append(val_loss)
    history["val_acc"].append(val_acc)

    print(f"Epoch {epoch+1}/{num_epochs}: "
          f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} "
          f"val_loss={val_loss:.4f} val_acc={val_acc:.4f}")

Example output:

Epoch 1/8: train_loss=0.4813 train_acc=0.8190 val_loss=0.1782 val_acc=1.0000
Epoch 2/8: train_loss=0.0172 train_acc=1.0000 val_loss=0.0076 val_acc=1.0000
Epoch 3/8: train_loss=0.0040 train_acc=1.0000 val_loss=0.0008 val_acc=1.0000
Epoch 4/8: train_loss=0.0014 train_acc=1.0000 val_loss=0.0002 val_acc=1.0000
Epoch 5/8: train_loss=0.0017 train_acc=1.0000 val_loss=0.0001 val_acc=1.0000
Epoch 6/8: train_loss=0.0013 train_acc=1.0000 val_loss=0.0001 val_acc=1.0000
Epoch 7/8: train_loss=0.0009 train_acc=1.0000 val_loss=0.0001 val_acc=1.0000
Epoch 8/8: train_loss=0.0021 train_acc=1.0000 val_loss=0.0000 val_acc=1.0000

Accuracy shoots up quickly in the first epoch and has essentially converged from the second epoch onward. Because this synthetic dataset has a relatively simple pattern, it converges quickly; with real data, convergence is often more gradual, and validation loss sometimes even turns upward partway through (a sign of overfitting). Visualizing these changes as a graph, rather than relying on numbers alone, helps you spot such shifts early.

import matplotlib.pyplot as plt

epochs_range = range(1, num_epochs + 1)
fig, axes = plt.subplots(1, 2, figsize=(10, 4))

axes[0].plot(epochs_range, history["train_loss"], label="Train Loss")
axes[0].plot(epochs_range, history["val_loss"], label="Val Loss")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("Loss")
axes[0].set_title("Loss Curve")
axes[0].legend()

axes[1].plot(epochs_range, history["train_acc"], label="Train Acc")
axes[1].plot(epochs_range, history["val_acc"], label="Val Acc")
axes[1].set_xlabel("Epoch")
axes[1].set_ylabel("Accuracy")
axes[1].set_title("Accuracy Curve")
axes[1].legend()

plt.tight_layout()
plt.savefig("training_curves.png")
print("Saved the training curves to training_curves.png")

Running this saves a graph to training_curves.png showing the loss curve dropping sharply toward zero and the accuracy curve climbing toward 1.0. If both the loss and accuracy curves improve smoothly, training is proceeding well. Conversely, if only the validation loss starts rising partway through, that's a sign of overfitting, and you should consider countermeasures such as strengthening data augmentation or reducing the number of epochs (early stopping).

5. Hyperparameter Tuning

A Hyperparameter is a value — such as the learning rate, batch size, or number of layers — that isn't determined automatically through training and must be set by a person beforehand (this is distinct from "parameters" such as weights, which are determined through training). Choosing appropriate hyperparameters can significantly change a model's performance and training stability.

Let's implement Grid Search — the simplest tuning method, which tries every combination of candidate values. To keep runtime manageable, we'll use a small number of epochs for each combination.

def build_and_train(lr, batch_size, num_epochs=2):
    tr_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
    va_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False)

    m = SimpleCNN(num_classes=3).to(device)
    opt = optim.Adam(m.parameters(), lr=lr)
    crit = nn.CrossEntropyLoss()

    for _ in range(num_epochs):
        train_one_epoch(m, tr_loader, crit, opt, device)

    _, val_acc = evaluate(m, va_loader, crit, device)
    return val_acc

results = {}
for lr in [0.1, 0.01, 0.001, 0.0001]:
    for batch_size in [16, 64]:
        acc = build_and_train(lr, batch_size, num_epochs=2)
        results[(lr, batch_size)] = acc
        print(f"lr={lr}, batch_size={batch_size} -> val_acc={acc:.4f}")

best_params = max(results, key=results.get)
print(f"\nBest combination: lr={best_params[0]}, batch_size={best_params[1]} "
      f"(val_acc={results[best_params]:.4f})")

Example output:

lr=0.1, batch_size=16 -> val_acc=1.0000
lr=0.1, batch_size=64 -> val_acc=1.0000
lr=0.01, batch_size=16 -> val_acc=1.0000
lr=0.01, batch_size=64 -> val_acc=1.0000
lr=0.001, batch_size=16 -> val_acc=1.0000
lr=0.001, batch_size=64 -> val_acc=1.0000
lr=0.0001, batch_size=16 -> val_acc=1.0000
lr=0.0001, batch_size=64 -> val_acc=0.9701

Best combination: lr=0.1, batch_size=16 (val_acc=1.0000)

You can see that only the combination of learning rate 0.0001 and batch size 64 has a slightly lower validation accuracy. When the learning rate stays small while the batch size grows, the number of weight updates per epoch decreases further, so with just 2 epochs the weights don't have enough time to fully converge. The other combinations converge quickly on this synthetic dataset, but with real data you'll typically see much clearer differences — for example, a learning rate that's too high causing the loss to diverge, or one that's too low causing slow convergence. Grid search's value lies precisely in surfacing these "dangerous combinations."

6. Model Deployment and Inference

Once training is complete, the model disappears from memory as soon as the process ends unless you save it. To reuse it, you need to save the parameters to a file so they can be loaded later. In PyTorch, the recommended approach is to save the state dictionary (state_dict) — a dictionary mapping each layer's name to its weight Tensor — rather than the entire model.

import os

save_path = "cnn_model.pth"
torch.save(model.state_dict(), save_path)
print(f"Saved the model to {save_path} "
      f"({os.path.getsize(save_path) / 1024:.1f} KB)")

Example output:

Saved the model to cnn_model.pth (418.7 KB)

To load it, first create a new SimpleCNN instance with the same structure, then load the saved weights into it. Specifying weights_only=True performs a safe load that doesn't include executable code.

loaded_model = SimpleCNN(num_classes=3).to(device)
loaded_model.load_state_dict(torch.load(save_path, map_location=device, weights_only=True))
loaded_model.eval()  # switch to inference mode (disables Dropout/BatchNorm)

# Run inference on a portion of the test data
sample_images, sample_labels = next(iter(test_loader))
sample_images = sample_images.to(device)

with torch.no_grad():
    logits = loaded_model(sample_images)
    probs = torch.softmax(logits, dim=1)
    preds = probs.argmax(dim=1)

print(f"Predicted classes: {preds[:8].cpu().tolist()}")
print(f"True classes: {sample_labels[:8].tolist()}")
print(f"Prediction confidence: {probs.max(dim=1).values[:8].cpu().tolist()}")

Example output:

Predicted classes: [2, 0, 1, 0, 2, 0, 0, 0]
True classes: [2, 0, 1, 0, 2, 0, 0, 0]
Prediction confidence: [0.9999..., 0.9999..., 0.9999..., 0.9999..., 0.9998..., 0.9999..., 0.9999..., 0.9999...]

There are three key points during inference.

• Always call model.eval(): disables Dropout and makes BatchNorm2d use its accumulated statistics instead of training-time behavior
• Stop gradient computation with torch.no_grad(): since backpropagation isn't needed during inference, this saves memory and computation time
• Convert to probabilities with Softmax: the model's raw output (logits) is hard to interpret directly, so converting it to per-class confidence with torch.softmax makes the results easier to explain

Revisiting the Learning Objectives

• ✅ Implement an image classification task end-to-end: implemented the full flow in code, from generating synthetic data to saving the model and running inference
• ✅ Build a CNN model: defined SimpleCNN, combining Conv2d, MaxPool2d, and Linear, and tracked how the feature map shape changes
• ✅ Apply data augmentation and batch normalization: combined random flips and rotations via transforms.Compose with training stabilization via nn.BatchNorm2d
• ✅ Visualize and monitor the training process: recorded per-epoch loss and accuracy and plotted training curves with matplotlib
• ✅ Perform basic hyperparameter tuning: compared combinations of learning rate and batch size via grid search and selected the best combination
• ✅ Run inference with a trained model: implemented saving and loading a state_dict and safe inference with model.eval() + torch.no_grad()

Exercises

class WiderCNN(nn.Module):
    def __init__(self, num_classes=3):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool = nn.MaxPool2d(2, 2)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)
        # 28 -> 14 -> 7, so the input dimension is 64 channels x 7 x 7
        self.fc1 = nn.Linear(64 * 7 * 7, 64)
        self.fc2 = nn.Linear(64, num_classes)

    def forward(self, x):
        x = self.pool(self.relu(self.bn1(self.conv1(x))))
        x = self.pool(self.relu(self.bn2(self.conv2(x))))
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

wider_model = WiderCNN(num_classes=3)
num_params = sum(p.numel() for p in wider_model.parameters())
print(f"Number of parameters: {num_params:,}")
# Example output: Number of parameters: 219,971 (a substantial increase over SimpleCNN)

def run_experiment(train_transform_to_use, num_epochs=5):
    train_ds = SyntheticImageDataset(all_images, all_labels, transform=train_transform_to_use)
    train_ds_subset = torch.utils.data.Subset(train_ds, train_subset.indices)
    loader = DataLoader(train_ds_subset, batch_size=32, shuffle=True)

    m = SimpleCNN(num_classes=3).to(device)
    opt = optim.Adam(m.parameters(), lr=0.001)
    crit = nn.CrossEntropyLoss()
    for _ in range(num_epochs):
        train_one_epoch(m, loader, crit, opt, device)
    _, test_acc = evaluate(m, test_loader, crit, device)
    return test_acc

acc_with_aug = run_experiment(train_transform)
acc_without_aug = run_experiment(eval_transform)
print(f"With augmentation: {acc_with_aug:.4f}, without augmentation: {acc_without_aug:.4f}")
# Example discussion: For simple data like this, where the pattern for each class is clearly
# separable, the difference between with and without augmentation tends to be small. On the
# other hand, the effect of augmentation becomes more pronounced under conditions such as
# similar patterns between classes, high noise, or a small amount of data.

def generate_synthetic_images_v2(num_samples_per_class=150, image_size=28,
                                   num_classes=4, seed=42):
    rng = np.random.default_rng(seed)
    images, labels = [], []
    for label in range(num_classes):
        for _ in range(num_samples_per_class):
            img = rng.normal(loc=0.15, scale=0.20,
                              size=(image_size, image_size)).astype(np.float32)
            if label == 0:
                img[8:12, :] += 0.35
                img[16:20, :] += 0.35
            elif label == 1:
                img[:, 8:12] += 0.35
                img[:, 16:20] += 0.35
            elif label == 2:
                idx = np.arange(image_size)
                img[idx, idx] += 0.35
                img[idx, image_size - 1 - idx] += 0.35
            else:  # grid pattern (4th class)
                img[::4, :] += 0.35
                img[:, ::4] += 0.35
            img = np.clip(img, 0.0, 1.0)
            images.append(img)
            labels.append(label)
    images = np.stack(images).astype(np.float32)[:, np.newaxis, :, :]
    labels = np.array(labels, dtype=np.int64)
    perm = rng.permutation(len(labels))
    return torch.from_numpy(images[perm]), torch.from_numpy(labels[perm])

images4, labels4 = generate_synthetic_images_v2(num_classes=4)
model4 = SimpleCNN(num_classes=4).to(device)
# From here, build the Dataset/DataLoader and run the training loop as in the main text

model4.eval()
with torch.no_grad():
    logits = model4(images4[:32].to(device))
    probs = torch.softmax(logits, dim=1)
    avg_confidence = probs.max(dim=1).values.mean().item()
print(f"Average prediction confidence: {avg_confidence:.4f}")

Summary

In this chapter, as the capstone of the PyTorch Basics Introduction series, we worked through a hands-on image classification project.

• ✅ Generated synthetic image data and split it into training, validation, and test sets using a custom Dataset and DataLoader
• ✅ Built a CNN model combining Conv2d, MaxPool2d, and Linear
• ✅ Combined data augmentation via transforms with training stabilization via BatchNorm2d
• ✅ Recorded the history of loss and accuracy and visualized it as training curves
• ✅ Compared hyperparameter combinations via grid search and selected the best configuration
• ✅ Saved and loaded a trained model and safely ran inference with model.eval() and torch.no_grad()

Reference Resources

• PyTorch Official Documentation
• torchvision Official Documentation
• PyTorch Official Tutorials
• PyTorch GitHub