Chapter 5: Image Classification Projects

Practical CNN Projects - From MNIST to CIFAR-10

📖 Reading Time: 30-35 minutes 📊 Difficulty: Intermediate to Advanced 💻 Code Examples: 17 📝 Exercises: 5

Learning Objectives

By reading this chapter, you will learn:


5.1 Project 1: MNIST Handwritten Digit Recognition

Dataset Preparation and Preprocessing

MNIST (Modified National Institute of Standards and Technology) is an image classification task for handwritten digits (0-9), often called the "Hello World" of machine learning.

Feature Details
Image Size 28×28 pixels (grayscale)
Number of Classes 10 classes (digits 0-9)
Training Data 60,000 images
Test Data 10,000 images
Difficulty Introductory level (best accuracy: 99.8%+)

Data Loading and Visualization

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

# Define data transforms
transform = transforms.Compose([
    transforms.ToTensor(),  # [0, 255] → [0.0, 1.0]
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST mean and standard deviation
])

# Load datasets
train_dataset = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

test_dataset = datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transform
)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

print(f'Training data: {len(train_dataset)} images')
print(f'Test data: {len(test_dataset)} images')

# Visualize sample images
def visualize_mnist_samples(dataset, n_samples=10):
    """Display MNIST sample images"""
    fig, axes = plt.subplots(2, 5, figsize=(12, 5))
    axes = axes.ravel()

    for i in range(n_samples):
        image, label = dataset[i]
        # Undo normalization
        image = image.squeeze() * 0.3081 + 0.1307

        axes[i].imshow(image, cmap='gray')
        axes[i].set_title(f'Label: {label}')
        axes[i].axis('off')

    plt.tight_layout()
    plt.show()

visualize_mnist_samples(train_dataset)

Designing the CNN Model

We design an efficient CNN architecture for MNIST.

graph LR A[Input
28×28×1] --> B[Conv1
24×24×32] B --> C[Pool1
12×12×32] C --> D[Conv2
8×8×64] D --> E[Pool2
4×4×64] E --> F[Flatten
1024] F --> G[FC1
128] G --> H[Dropout
0.5] H --> I[FC2
10] style A fill:#e3f2fd style I fill:#e8f5e9
class MNISTNet(nn.Module):
    """CNN model for MNIST"""

    def __init__(self):
        super(MNISTNet, self).__init__()

        # Convolutional layers
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5)  # 28×28×1 → 24×24×32
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5)  # 12×12×32 → 8×8×64

        # Pooling layer
        self.pool = nn.MaxPool2d(2, 2)

        # Fully connected layers
        self.fc1 = nn.Linear(64 * 4 * 4, 128)
        self.fc2 = nn.Linear(128, 10)

        # Dropout
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        # Block 1: Conv → ReLU → Pool
        x = self.pool(F.relu(self.conv1(x)))  # → 12×12×32

        # Block 2: Conv → ReLU → Pool
        x = self.pool(F.relu(self.conv2(x)))  # → 4×4×64

        # Flatten
        x = x.view(-1, 64 * 4 * 4)  # → 1024

        # Fully connected layers
        x = F.relu(self.fc1(x))  # → 128
        x = self.dropout(x)
        x = self.fc2(x)  # → 10

        return x

# Instantiate the model
model = MNISTNet().to(device)

# Display model structure
print(model)

# Count parameters
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'\nTotal parameters: {total_params:,}')
print(f'Trainable parameters: {trainable_params:,}')

Example output:

MNISTNet(
  (conv1): Conv2d(1, 32, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (fc1): Linear(in_features=1024, out_features=128, bias=True)
  (fc2): Linear(in_features=128, out_features=10, bias=True)
  (dropout): Dropout(p=0.5, inplace=False)
)

Total parameters: 163,978
Trainable parameters: 163,978

Training and Evaluation

def train_epoch(model, device, train_loader, optimizer, criterion, epoch):
    """Train for one epoch"""
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        # Zero the gradients
        optimizer.zero_grad()

        # Forward pass
        output = model(data)
        loss = criterion(output, target)

        # Backward pass
        loss.backward()
        optimizer.step()

        # Statistics
        running_loss += loss.item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()

        # Progress display
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch}, Batch {batch_idx}/{len(train_loader)}, '
                  f'Loss: {loss.item():.4f}, Acc: {100.*correct/total:.2f}%')

    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100. * correct / total

    return epoch_loss, epoch_acc

def evaluate(model, device, test_loader, criterion):
    """Evaluate on test data"""
    model.eval()
    test_loss = 0
    correct = 0

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader)
    test_acc = 100. * correct / len(test_loader.dataset)

    print(f'\nTest set: Average loss: {test_loss:.4f}, '
          f'Accuracy: {correct}/{len(test_loader.dataset)} ({test_acc:.2f}%)\n')

    return test_loss, test_acc

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

# Training loop
num_epochs = 10
train_losses, train_accs = [], []
test_losses, test_accs = [], []

for epoch in range(1, num_epochs + 1):
    train_loss, train_acc = train_epoch(model, device, train_loader, optimizer, criterion, epoch)
    test_loss, test_acc = evaluate(model, device, test_loader, criterion)

    train_losses.append(train_loss)
    train_accs.append(train_acc)
    test_losses.append(test_loss)
    test_accs.append(test_acc)

# Save the model
torch.save(model.state_dict(), 'mnist_cnn.pth')
print('Model saved: mnist_cnn.pth')

Error Analysis and Visualization

from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns

def plot_confusion_matrix(model, device, test_loader):
    """Visualize the confusion matrix"""
    model.eval()
    all_preds = []
    all_targets = []

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1)
            all_preds.extend(pred.cpu().numpy())
            all_targets.extend(target.cpu().numpy())

    # Compute the confusion matrix
    cm = confusion_matrix(all_targets, all_preds)

    # Visualization
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=range(10), yticklabels=range(10))
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.title('Confusion Matrix - MNIST')
    plt.show()

    # Classification report
    print('\nClassification Report:')
    print(classification_report(all_targets, all_preds,
                                target_names=[str(i) for i in range(10)]))

plot_confusion_matrix(model, device, test_loader)

def visualize_misclassified(model, device, test_loader, n_samples=10):
    """Display misclassified samples"""
    model.eval()
    misclassified = []

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1)

            # Find misclassifications
            mask = pred != target
            if mask.sum() > 0:
                for i in range(len(mask)):
                    if mask[i]:
                        misclassified.append({
                            'image': data[i].cpu(),
                            'true': target[i].item(),
                            'pred': pred[i].item(),
                            'confidence': F.softmax(output[i], dim=0)[pred[i]].item()
                        })

                        if len(misclassified) >= n_samples:
                            break

            if len(misclassified) >= n_samples:
                break

    # Visualization
    fig, axes = plt.subplots(2, 5, figsize=(12, 5))
    axes = axes.ravel()

    for i, item in enumerate(misclassified[:n_samples]):
        image = item['image'].squeeze() * 0.3081 + 0.1307
        axes[i].imshow(image, cmap='gray')
        axes[i].set_title(f"True: {item['true']}, Pred: {item['pred']}\n"
                         f"Conf: {item['confidence']:.2%}")
        axes[i].axis('off')

    plt.tight_layout()
    plt.show()

visualize_misclassified(model, device, test_loader)

def plot_training_history(train_losses, train_accs, test_losses, test_accs):
    """Visualize learning curves"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    epochs = range(1, len(train_losses) + 1)

    # Loss curves
    ax1.plot(epochs, train_losses, 'b-', label='Training Loss', linewidth=2)
    ax1.plot(epochs, test_losses, 'r-', label='Test Loss', linewidth=2)
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    ax1.set_title('Training and Test Loss')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # Accuracy curves
    ax2.plot(epochs, train_accs, 'b-', label='Training Accuracy', linewidth=2)
    ax2.plot(epochs, test_accs, 'r-', label='Test Accuracy', linewidth=2)
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Accuracy (%)')
    ax2.set_title('Training and Test Accuracy')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()

plot_training_history(train_losses, train_accs, test_losses, test_accs)

5.2 Project 2: CIFAR-10 Color Image Classification

Dataset Overview

CIFAR-10 (Canadian Institute for Advanced Research) is a real-world image classification task that is more challenging than MNIST.

Feature Details
Image Size 32×32 pixels (RGB color)
Number of Classes 10 classes (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck)
Training Data 50,000 images (5,000 per class)
Test Data 10,000 images (1,000 per class)
Difficulty Intermediate (best accuracy: 99%+, typically 90-95%)

The Importance of Data Augmentation

For CIFAR-10, data augmentation is the key to improving accuracy.

graph LR A[Original Image] --> B[Horizontal Flip] A --> C[Random Crop] A --> D[Color Jitter] A --> E[Rotation] B --> F[Training Data] C --> F D --> F E --> F style A fill:#e3f2fd style F fill:#e8f5e9
from torchvision import datasets, transforms

# CIFAR-10 class names
classes = ('airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

# Training transforms including data augmentation
train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),  # Random crop
    transforms.RandomHorizontalFlip(),  # Horizontal flip (50% probability)
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),  # Color jitter
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])

# Test transforms (no augmentation)
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])

# Load datasets
train_dataset = datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=train_transform
)

test_dataset = datasets.CIFAR10(
    root='./data',
    train=False,
    download=True,
    transform=test_transform
)

# Data loaders
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

def visualize_augmentation(dataset, idx=0, n_augments=8):
    """Visualize the effect of data augmentation"""
    fig, axes = plt.subplots(2, 4, figsize=(12, 6))
    axes = axes.ravel()

    # Original image
    original_img, label = dataset[idx]

    for i in range(n_augments):
        # Get an augmented image
        img, _ = dataset[idx]

        # Undo normalization
        img = img.numpy().transpose((1, 2, 0))
        mean = np.array([0.4914, 0.4822, 0.4465])
        std = np.array([0.2470, 0.2435, 0.2616])
        img = std * img + mean
        img = np.clip(img, 0, 1)

        axes[i].imshow(img)
        axes[i].set_title(f'Augmentation {i+1}')
        axes[i].axis('off')

    plt.suptitle(f'Class: {classes[label]}', fontsize=16)
    plt.tight_layout()
    plt.show()

visualize_augmentation(train_dataset)

Designing a Deeper Network

A deeper VGG-style CNN is effective for CIFAR-10.

class CIFAR10Net(nn.Module):
    """Deep CNN for CIFAR-10 (VGG-style)"""

    def __init__(self, num_classes=10):
        super(CIFAR10Net, self).__init__()

        # Block 1: 32×32 → 16×16
        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.bn1_1 = nn.BatchNorm2d(64)
        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.bn1_2 = nn.BatchNorm2d(64)

        # Block 2: 16×16 → 8×8
        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2_1 = nn.BatchNorm2d(128)
        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.bn2_2 = nn.BatchNorm2d(128)

        # Block 3: 8×8 → 4×4
        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.bn3_1 = nn.BatchNorm2d(256)
        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.bn3_2 = nn.BatchNorm2d(256)
        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.bn3_3 = nn.BatchNorm2d(256)

        # Pooling
        self.pool = nn.MaxPool2d(2, 2)

        # Fully connected layers
        self.fc1 = nn.Linear(256 * 4 * 4, 512)
        self.bn_fc1 = nn.BatchNorm1d(512)
        self.fc2 = nn.Linear(512, num_classes)

        # Dropout
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        # Block 1
        x = F.relu(self.bn1_1(self.conv1_1(x)))
        x = F.relu(self.bn1_2(self.conv1_2(x)))
        x = self.pool(x)

        # Block 2
        x = F.relu(self.bn2_1(self.conv2_1(x)))
        x = F.relu(self.bn2_2(self.conv2_2(x)))
        x = self.pool(x)

        # Block 3
        x = F.relu(self.bn3_1(self.conv3_1(x)))
        x = F.relu(self.bn3_2(self.conv3_2(x)))
        x = F.relu(self.bn3_3(self.conv3_3(x)))
        x = self.pool(x)

        # Flatten
        x = x.view(-1, 256 * 4 * 4)

        # Fully connected layers
        x = F.relu(self.bn_fc1(self.fc1(x)))
        x = self.dropout(x)
        x = self.fc2(x)

        return x

# Instantiate the model
model = CIFAR10Net().to(device)

# Parameter count
total_params = sum(p.numel() for p in model.parameters())
print(f'Total parameters: {total_params:,}')

Regularization Techniques (Dropout, Batch Normalization)

Batch Normalization: Normalizes the output of each layer, stabilizing and accelerating training.

Dropout: Randomly deactivates neurons during training to prevent overfitting.

def train_with_scheduler(model, device, train_loader, test_loader, num_epochs=50):
    """Training with a Learning Rate Scheduler"""

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)

    # Learning Rate Scheduler
    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[25, 40], gamma=0.1)

    train_losses, train_accs = [], []
    test_losses, test_accs = [], []

    best_acc = 0.0

    for epoch in range(1, num_epochs + 1):
        # Training
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()

        train_loss = running_loss / len(train_loader)
        train_acc = 100. * correct / total

        # Evaluation
        model.eval()
        test_loss = 0
        correct = 0

        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                test_loss += criterion(output, target).item()
                pred = output.argmax(dim=1)
                correct += pred.eq(target).sum().item()

        test_loss /= len(test_loader)
        test_acc = 100. * correct / len(test_loader.dataset)

        # Scheduler step
        scheduler.step()

        # Record
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        test_losses.append(test_loss)
        test_accs.append(test_acc)

        print(f'Epoch {epoch}/{num_epochs} - '
              f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% - '
              f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}% - '
              f'LR: {scheduler.get_last_lr()[0]:.6f}')

        # Save the best model
        if test_acc > best_acc:
            best_acc = test_acc
            torch.save(model.state_dict(), 'cifar10_best.pth')
            print(f'  → Best model saved! (Acc: {best_acc:.2f}%)')

    return train_losses, train_accs, test_losses, test_accs

# Run training
train_losses, train_accs, test_losses, test_accs = train_with_scheduler(
    model, device, train_loader, test_loader, num_epochs=50
)

Early Stopping Implementation

class EarlyStopping:
    """Early Stopping implementation"""

    def __init__(self, patience=7, min_delta=0, path='checkpoint.pth'):
        """
        Args:
            patience: Maximum number of epochs without improvement
            min_delta: Minimum change to be considered an improvement
            path: Path to save the model
        """
        self.patience = patience
        self.min_delta = min_delta
        self.path = path
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.best_loss = np.Inf

    def __call__(self, val_loss, model):
        score = -val_loss

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
        elif score < self.best_score + self.min_delta:
            self.counter += 1
            print(f'EarlyStopping counter: {self.counter}/{self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
            self.counter = 0

    def save_checkpoint(self, val_loss, model):
        """Save the model"""
        torch.save(model.state_dict(), self.path)
        self.best_loss = val_loss

# Usage example
def train_with_early_stopping(model, device, train_loader, test_loader,
                              num_epochs=100, patience=10):
    """Training with Early Stopping"""

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    early_stopping = EarlyStopping(patience=patience, path='cifar10_checkpoint.pth')

    for epoch in range(1, num_epochs + 1):
        # Training loop (omitted)
        train_loss = 0.0  # Actually run training here

        # Validation
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                val_loss += criterion(output, target).item()

        val_loss /= len(test_loader)

        # Early Stopping check
        early_stopping(val_loss, model)

        if early_stopping.early_stop:
            print(f'Early stopping triggered at epoch {epoch}')
            break

    # Load the best model
    model.load_state_dict(torch.load('cifar10_checkpoint.pth'))
    return model

5.3 Advanced Techniques

Transfer Learning

By using pretrained models, you can achieve high accuracy even with limited data.

graph LR A[ImageNet
Pretraining] --> B[Feature Extractor
Frozen] B --> C[Classifier
Training] C --> D[CIFAR-10
Classification] style A fill:#e3f2fd style B fill:#fff3e0 style C fill:#f3e5f5 style D fill:#e8f5e9
import torchvision.models as models

def create_transfer_model(num_classes=10, freeze_features=True):
    """Transfer learning model based on ResNet18"""

    # Load ResNet18 pretrained on ImageNet
    model = models.resnet18(pretrained=True)

    # Freeze the feature extractor
    if freeze_features:
        for param in model.parameters():
            param.requires_grad = False

    # Replace the final layer (10-class classification for CIFAR-10)
    num_features = model.fc.in_features
    model.fc = nn.Linear(num_features, num_classes)

    return model

# Create the model
transfer_model = create_transfer_model(num_classes=10, freeze_features=True).to(device)

# Optimize only the trainable parameters
optimizer = optim.Adam(filter(lambda p: p.requires_grad, transfer_model.parameters()),
                       lr=0.001)

print("Transfer learning model:")
print(f"Total parameters: {sum(p.numel() for p in transfer_model.parameters()):,}")
print(f"Trainable parameters: {sum(p.numel() for p in transfer_model.parameters() if p.requires_grad):,}")

def finetune_transfer_model(model, device, train_loader, test_loader, num_epochs=20):
    """Fine-tuning the transfer learning model"""

    criterion = nn.CrossEntropyLoss()

    # Phase 1: Train the classifier only (5 epochs)
    print("\nPhase 1: Training classifier only...")
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)

    for epoch in range(1, 6):
        model.train()
        running_loss = 0.0

        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f"Epoch {epoch}/5 - Loss: {running_loss/len(train_loader):.4f}")

    # Phase 2: Fine-tune all layers (15 epochs)
    print("\nPhase 2: Fine-tuning all layers...")
    for param in model.parameters():
        param.requires_grad = True

    optimizer = optim.Adam(model.parameters(), lr=0.0001)  # Smaller learning rate

    for epoch in range(1, 16):
        model.train()
        running_loss = 0.0

        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        # Evaluation
        model.eval()
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                pred = output.argmax(dim=1)
                correct += pred.eq(target).sum().item()

        test_acc = 100. * correct / len(test_loader.dataset)
        print(f"Epoch {epoch}/15 - Loss: {running_loss/len(train_loader):.4f}, "
              f"Test Acc: {test_acc:.2f}%")

    return model

# Run fine-tuning
transfer_model = finetune_transfer_model(transfer_model, device, train_loader, test_loader)

Learning Rate Finder

A technique for automatically finding the optimal learning rate (made famous by Fast.ai).

import matplotlib.pyplot as plt

class LRFinder:
    """Learning Rate Finder implementation"""

    def __init__(self, model, optimizer, criterion, device):
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.device = device
        self.history = {'lr': [], 'loss': []}
        self.best_loss = 1e9

    def range_test(self, train_loader, start_lr=1e-7, end_lr=10, num_iter=100):
        """Test a range of learning rates"""

        # Learning rate schedule
        lr_schedule = np.logspace(np.log10(start_lr), np.log10(end_lr), num_iter)

        # Save the initial model state
        model_state = self.model.state_dict()
        optimizer_state = self.optimizer.state_dict()

        self.model.train()
        iter_count = 0

        for batch_idx, (data, target) in enumerate(train_loader):
            if iter_count >= num_iter:
                break

            data, target = data.to(self.device), target.to(self.device)

            # Update the learning rate
            lr = lr_schedule[iter_count]
            for param_group in self.optimizer.param_groups:
                param_group['lr'] = lr

            # Forward and backward pass
            self.optimizer.zero_grad()
            output = self.model(data)
            loss = self.criterion(output, target)
            loss.backward()
            self.optimizer.step()

            # Record
            self.history['lr'].append(lr)
            self.history['loss'].append(loss.item())

            # Divergence check
            if loss.item() > 4 * self.best_loss:
                break

            if loss.item() < self.best_loss:
                self.best_loss = loss.item()

            iter_count += 1

        # Restore the model to its initial state
        self.model.load_state_dict(model_state)
        self.optimizer.load_state_dict(optimizer_state)

    def plot(self):
        """Plot the results"""
        plt.figure(figsize=(10, 6))
        plt.plot(self.history['lr'], self.history['loss'])
        plt.xscale('log')
        plt.xlabel('Learning Rate')
        plt.ylabel('Loss')
        plt.title('Learning Rate Finder')
        plt.grid(True, alpha=0.3)
        plt.show()

        # Suggested learning rate
        min_loss_idx = np.argmin(self.history['loss'])
        suggested_lr = self.history['lr'][min_loss_idx] / 10
        print(f"Suggested learning rate: {suggested_lr:.2e}")

# Usage example
model = CIFAR10Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
criterion = nn.CrossEntropyLoss()

lr_finder = LRFinder(model, optimizer, criterion, device)
lr_finder.range_test(train_loader, start_lr=1e-6, end_lr=1, num_iter=100)
lr_finder.plot()

Model Ensembles

Combining predictions from multiple models improves accuracy.

class ModelEnsemble:
    """Model ensemble implementation"""

    def __init__(self, models, device):
        """
        Args:
            models: List of models
            device: Execution device
        """
        self.models = models
        self.device = device

        # Set all models to evaluation mode
        for model in self.models:
            model.eval()

    def predict(self, data):
        """Ensemble prediction (averaging)"""
        predictions = []

        with torch.no_grad():
            for model in self.models:
                output = model(data)
                predictions.append(F.softmax(output, dim=1))

        # Take the average
        ensemble_pred = torch.stack(predictions).mean(dim=0)
        return ensemble_pred

    def evaluate(self, test_loader):
        """Evaluate on test data"""
        correct = 0
        total = 0

        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(self.device), target.to(self.device)

                # Ensemble prediction
                ensemble_pred = self.predict(data)
                pred = ensemble_pred.argmax(dim=1)

                correct += pred.eq(target).sum().item()
                total += target.size(0)

        accuracy = 100. * correct / total
        print(f'Ensemble Accuracy: {correct}/{total} ({accuracy:.2f}%)')
        return accuracy

# Train multiple models (with different initializations or settings)
def train_multiple_models(n_models=5):
    """Train multiple models"""
    models = []

    for i in range(n_models):
        print(f"\nTraining model {i+1}/{n_models}...")

        # Model initialization
        model = CIFAR10Net().to(device)

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

        # Training loop (omitted)
        # ... train_epoch(model, device, train_loader, optimizer, criterion)

        models.append(model)

    return models

# Ensemble usage example
# models = train_multiple_models(n_models=5)
# ensemble = ModelEnsemble(models, device)
# ensemble_acc = ensemble.evaluate(test_loader)

Grad-CAM Visualization

Visualize where in the image the model is paying attention.

class GradCAM:
    """Gradient-weighted Class Activation Mapping implementation"""

    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None

        # Register hooks
        target_layer.register_forward_hook(self.save_activation)
        target_layer.register_backward_hook(self.save_gradient)

    def save_activation(self, module, input, output):
        """Save activations during the forward pass"""
        self.activations = output.detach()

    def save_gradient(self, module, grad_input, grad_output):
        """Save gradients during the backward pass"""
        self.gradients = grad_output[0].detach()

    def generate_cam(self, input_image, target_class):
        """Generate the CAM"""

        # Forward pass
        output = self.model(input_image)

        # Gradient with respect to the target class score
        self.model.zero_grad()
        class_loss = output[0, target_class]
        class_loss.backward()

        # Get gradients and activations
        gradients = self.gradients[0]  # [C, H, W]
        activations = self.activations[0]  # [C, H, W]

        # Average of gradients (weights)
        weights = gradients.mean(dim=(1, 2))  # [C]

        # Weighted sum
        cam = torch.zeros(activations.shape[1:], dtype=torch.float32)
        for i, w in enumerate(weights):
            cam += w * activations[i]

        # Apply ReLU
        cam = F.relu(cam)

        # Normalize
        cam = cam - cam.min()
        cam = cam / cam.max()

        return cam.cpu().numpy()

def visualize_gradcam(model, image, label, device):
    """Visualize Grad-CAM"""

    # Prepare Grad-CAM (target the last convolutional layer)
    target_layer = model.conv3_3  # For CIFAR10Net
    gradcam = GradCAM(model, target_layer)

    # Generate the CAM
    model.eval()
    image_input = image.unsqueeze(0).to(device)
    cam = gradcam.generate_cam(image_input, label)

    # Prepare the original image
    img = image.cpu().numpy().transpose((1, 2, 0))
    mean = np.array([0.4914, 0.4822, 0.4465])
    std = np.array([0.2470, 0.2435, 0.2616])
    img = std * img + mean
    img = np.clip(img, 0, 1)

    # Resize the CAM
    from scipy.ndimage import zoom
    cam_resized = zoom(cam, (img.shape[0]/cam.shape[0], img.shape[1]/cam.shape[1]))

    # Visualization
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))

    axes[0].imshow(img)
    axes[0].set_title('Original Image')
    axes[0].axis('off')

    axes[1].imshow(cam_resized, cmap='jet')
    axes[1].set_title('Grad-CAM')
    axes[1].axis('off')

    axes[2].imshow(img)
    axes[2].imshow(cam_resized, cmap='jet', alpha=0.5)
    axes[2].set_title('Overlay')
    axes[2].axis('off')

    plt.tight_layout()
    plt.show()

# Usage example
# image, label = test_dataset[0]
# visualize_gradcam(model, image, label, device)

5.4 Practical Tips and Best Practices

Hyperparameter Tuning

graph TD A[Initial Settings] --> B{Validation Accuracy} B -->|Low| C[Adjust Learning Rate] B -->|Overfitting| D[Strengthen Regularization] B -->|Underfitting| E[Increase Epochs] C --> F[Retrain] D --> F E --> F F --> B B -->|Satisfied| G[Final Evaluation] style A fill:#e3f2fd style G fill:#e8f5e9
Hyperparameter Recommended Range Tuning Tips
Learning Rate 1e-4 ~ 1e-1 Use LR Finder; too large causes divergence
Batch Size 32 ~ 256 Depends on GPU memory; larger is more stable
Weight Decay 1e-5 ~ 1e-3 Counteracts overfitting; L2 regularization
Dropout Rate 0.3 ~ 0.5 Increase if overfitting is severe
Optimizer Adam, SGD+Momentum Adam is general-purpose; SGD converges well

Debugging Techniques

class ModelDebugger:
    """Model debugging tool"""

    @staticmethod
    def check_gradients(model):
        """Check gradients"""
        print("\n=== Gradient Check ===")
        for name, param in model.named_parameters():
            if param.requires_grad and param.grad is not None:
                grad_mean = param.grad.mean().item()
                grad_std = param.grad.std().item()
                grad_max = param.grad.abs().max().item()
                print(f"{name:30s} - Mean: {grad_mean:8.6f}, "
                      f"Std: {grad_std:8.6f}, Max: {grad_max:8.6f}")

                # Vanishing/exploding gradient warnings
                if grad_max < 1e-6:
                    print(f"  ⚠️  WARNING: Vanishing gradient!")
                if grad_max > 100:
                    print(f"  ⚠️  WARNING: Exploding gradient!")

    @staticmethod
    def check_weights(model):
        """Check weight statistics"""
        print("\n=== Weight Statistics ===")
        for name, param in model.named_parameters():
            if 'weight' in name:
                weight_mean = param.data.mean().item()
                weight_std = param.data.std().item()
                print(f"{name:30s} - Mean: {weight_mean:8.6f}, Std: {weight_std:8.6f}")

    @staticmethod
    def check_nan_inf(model):
        """Detect NaN/Inf"""
        print("\n=== NaN/Inf Check ===")
        has_issue = False
        for name, param in model.named_parameters():
            if torch.isnan(param.data).any():
                print(f"  ❌ NaN detected in {name}")
                has_issue = True
            if torch.isinf(param.data).any():
                print(f"  ❌ Inf detected in {name}")
                has_issue = True

        if not has_issue:
            print("  ✅ No NaN/Inf detected")

    @staticmethod
    def visualize_activation_distribution(model, data, device):
        """Visualize the distribution of activations"""
        activations = {}

        def hook_fn(name):
            def hook(module, input, output):
                activations[name] = output.detach().cpu()
            return hook

        # Register hooks
        hooks = []
        for name, module in model.named_modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                hooks.append(module.register_forward_hook(hook_fn(name)))

        # Forward pass
        model.eval()
        with torch.no_grad():
            _ = model(data.to(device))

        # Remove hooks
        for hook in hooks:
            hook.remove()

        # Visualization
        fig, axes = plt.subplots(len(activations), 1, figsize=(10, 3*len(activations)))
        if len(activations) == 1:
            axes = [axes]

        for ax, (name, activation) in zip(axes, activations.items()):
            activation_flat = activation.flatten().numpy()
            ax.hist(activation_flat, bins=50, alpha=0.7)
            ax.set_title(f'Activation Distribution: {name}')
            ax.set_xlabel('Activation Value')
            ax.set_ylabel('Frequency')
            ax.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.show()

# Usage example
debugger = ModelDebugger()

# Use inside the training loop
for epoch in range(num_epochs):
    # Training
    # ...

    # Debug checks
    debugger.check_gradients(model)
    debugger.check_nan_inf(model)

GPU Utilization and Memory Optimization

import torch.cuda as cuda

class GPUOptimizer:
    """GPU optimization utilities"""

    @staticmethod
    def get_gpu_info():
        """Get GPU information"""
        if not torch.cuda.is_available():
            print("CUDA is not available")
            return

        print(f"GPU Device: {torch.cuda.get_device_name(0)}")
        print(f"CUDA Version: {torch.version.cuda}")
        print(f"Total Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
        print(f"Allocated Memory: {torch.cuda.memory_allocated(0) / 1e9:.4f} GB")
        print(f"Cached Memory: {torch.cuda.memory_reserved(0) / 1e9:.4f} GB")

    @staticmethod
    def clear_cache():
        """Clear the GPU cache"""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            print("GPU cache cleared")

    @staticmethod
    def mixed_precision_training_example(model, train_loader, device):
        """Mixed Precision Training (FP16) example"""
        from torch.cuda.amp import autocast, GradScaler

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

        model.train()

        for data, target in train_loader:
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()

            # Mixed Precision
            with autocast():
                output = model(data)
                loss = criterion(output, target)

            # Backpropagation using the scaler
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

        print("Mixed precision training completed")

# Display GPU information
GPUOptimizer.get_gpu_info()

# Memory optimization tips
def optimize_memory_usage():
    """Optimize memory usage"""

    tips = [
        "1. Reduce the batch size (64 → 32 → 16)",
        "2. Use Gradient Accumulation (accumulate gradients over multiple steps)",
        "3. Use Mixed Precision Training (FP16)",
        "4. Wrap unnecessary intermediate results in torch.no_grad()",
        "5. Explicitly delete unneeded variables with del",
        "6. Clear the cache with torch.cuda.empty_cache()",
        "7. Adjust DataLoader num_workers (CPU/GPU balance)",
        "8. Use inplace operations (x = x + 1 → x += 1)"
    ]

    print("\n=== Memory Optimization Tips ===")
    for tip in tips:
        print(tip)

optimize_memory_usage()

Deploying to Production

import torch.jit

class ModelDeployment:
    """Model deployment"""

    @staticmethod
    def export_to_torchscript(model, example_input, save_path='model_scripted.pt'):
        """Export to TorchScript"""
        model.eval()

        # Script mode
        scripted_model = torch.jit.script(model)
        scripted_model.save(save_path)
        print(f"Model exported to TorchScript: {save_path}")

        return scripted_model

    @staticmethod
    def export_to_onnx(model, example_input, save_path='model.onnx'):
        """Export to ONNX format"""
        model.eval()

        torch.onnx.export(
            model,
            example_input,
            save_path,
            export_params=True,
            opset_version=11,
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        )
        print(f"Model exported to ONNX: {save_path}")

    @staticmethod
    def optimize_for_inference(model):
        """Optimization for inference"""
        model.eval()

        # Set to inference mode
        for param in model.parameters():
            param.requires_grad = False

        # Freeze BatchNorm and Dropout
        for module in model.modules():
            if isinstance(module, nn.BatchNorm2d):
                module.track_running_stats = False
            if isinstance(module, nn.Dropout):
                module.p = 0

        return model

    @staticmethod
    def benchmark_inference(model, example_input, device, num_runs=100):
        """Benchmark inference speed"""
        import time

        model.eval()
        model = model.to(device)
        example_input = example_input.to(device)

        # Warm-up
        with torch.no_grad():
            for _ in range(10):
                _ = model(example_input)

        # Benchmark
        torch.cuda.synchronize() if device.type == 'cuda' else None
        start_time = time.time()

        with torch.no_grad():
            for _ in range(num_runs):
                _ = model(example_input)

        torch.cuda.synchronize() if device.type == 'cuda' else None
        end_time = time.time()

        avg_time = (end_time - start_time) / num_runs
        fps = 1 / avg_time

        print(f"\n=== Inference Benchmark ===")
        print(f"Average inference time: {avg_time*1000:.2f} ms")
        print(f"Throughput: {fps:.2f} FPS")
        print(f"Total runs: {num_runs}")

# Usage example
deployment = ModelDeployment()

# Load the model
model = CIFAR10Net().to(device)
model.load_state_dict(torch.load('cifar10_best.pth'))

# Optimize for inference
model = deployment.optimize_for_inference(model)

# Export
example_input = torch.randn(1, 3, 32, 32).to(device)
# deployment.export_to_torchscript(model, example_input)
# deployment.export_to_onnx(model, example_input)

# Benchmark
deployment.benchmark_inference(model, example_input, device)

Chapter Summary

What We Learned

  1. MNIST Handwritten Digit Recognition

    • A complete data preparation pipeline
    • Efficient CNN architecture design
    • Confusion matrices and error analysis
  2. CIFAR-10 Color Image Classification

    • Practical use of data augmentation
    • Deep VGG-style networks
    • Batch Normalization and Dropout
  3. Advanced Techniques

    • Transfer learning and fine-tuning
    • Learning Rate Finder
    • Model ensembles
    • Grad-CAM visualization
  4. Practical Skills

    • Hyperparameter tuning
    • Debugging techniques
    • GPU optimization
    • Production deployment

Best Practices Checklist

Item Importance Description
✅ Data Normalization Essential Normalize to mean 0, standard deviation 1
✅ Data Augmentation High Increase the diversity of training data
✅ Batch Normalization High Stabilize and speed up training
✅ Dropout High Prevent overfitting
✅ Learning Rate Scheduling Medium Dynamic learning rate adjustment
✅ Early Stopping Medium Early detection of overfitting
✅ Model Saving Essential Checkpoints of the best model
✅ Recording Evaluation Metrics Essential Visualization of the training process

Exercises

Exercise 1 (Difficulty: medium)

Apply the following improvements to the MNIST model to reach an accuracy of 99% or higher:

Hint
Sample Solution
class ImprovedMNISTNet(nn.Module):
    """Improved MNIST network"""

    def __init__(self):
        super(ImprovedMNISTNet, self).__init__()

        # Convolutional layers (3 blocks)
        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.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)

        self.pool = nn.MaxPool2d(2, 2)

        # Fully connected layers
        self.fc1 = nn.Linear(128 * 3 * 3, 256)
        self.bn_fc = nn.BatchNorm1d(256)
        self.fc2 = nn.Linear(256, 10)

        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        # Block 1: 28×28 → 14×14
        x = self.pool(F.relu(self.bn1(self.conv1(x))))

        # Block 2: 14×14 → 7×7
        x = self.pool(F.relu(self.bn2(self.conv2(x))))

        # Block 3: 7×7 → 3×3
        x = self.pool(F.relu(self.bn3(self.conv3(x))))

        # Flatten and fully connected layers
        x = x.view(-1, 128 * 3 * 3)
        x = F.relu(self.bn_fc(self.fc1(x)))
        x = self.dropout(x)
        x = self.fc2(x)

        return x

# Data augmentation
train_transform = transforms.Compose([
    transforms.RandomRotation(10),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# Training
# model = ImprovedMNISTNet().to(device)
# Expected accuracy: 99.2% or higher

Exercise 2 (Difficulty: hard)

Using transfer learning on CIFAR-10, achieve 90% or higher accuracy within 20 epochs. Use ResNet50 and implement an appropriate fine-tuning strategy.

Hint

Exercise 3 (Difficulty: medium)

Implement a Learning Rate Finder and find the optimal learning rate. Train the model using that learning rate and compare the results.

Sample Solution
# Use the LRFinder
model = CIFAR10Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

lr_finder = LRFinder(model, optimizer, criterion, device)
lr_finder.range_test(train_loader, start_lr=1e-6, end_lr=1, num_iter=200)
lr_finder.plot()

# Use the suggested learning rate
# Example output: "Suggested learning rate: 1.2e-02"
suggested_lr = 0.012

# Train with a new model
model = CIFAR10Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=suggested_lr, momentum=0.9)
# Run training...

Exercise 4 (Difficulty: hard)

Train three different models (different architectures or different initializations) and improve accuracy with an ensemble. Aim for an improvement of 2% or more over the individual models.

Hint

Exercise 5 (Difficulty: hard)

Implement Mixed Precision Training (FP16) and compare it with regular training. Measure the differences in speed, memory usage, and accuracy.

Sample Solution
from torch.cuda.amp import autocast, GradScaler
import time

def train_with_mixed_precision(model, train_loader, num_epochs=10):
    """Mixed Precision Training implementation"""

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

    start_time = time.time()

    for epoch in range(1, num_epochs + 1):
        model.train()
        running_loss = 0.0

        for data, target in train_loader:
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()

            # Mixed Precision
            with autocast():
                output = model(data)
                loss = criterion(output, target)

            # Scaled backpropagation
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            running_loss += loss.item()

        epoch_loss = running_loss / len(train_loader)
        print(f"Epoch {epoch}/{num_epochs} - Loss: {epoch_loss:.4f}")

    total_time = time.time() - start_time
    print(f"\nTotal training time: {total_time:.2f} seconds")

    # Memory usage
    if torch.cuda.is_available():
        print(f"Max memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

# Compare with regular training
model_fp32 = CIFAR10Net().to(device)
model_fp16 = CIFAR10Net().to(device)

print("=== FP32 Training ===")
# Regular training...

print("\n=== FP16 Training ===")
train_with_mixed_precision(model_fp16, train_loader)

# Expected results:
# - FP16 is 1.5-2x faster
# - Memory usage reduced by 30-40%
# - Accuracy nearly identical (within ±0.5%)

References

  1. LeCun, Y., Bottou, L., Bengio, Y., & Haffner, P. (1998). "Gradient-based learning applied to document recognition." Proceedings of the IEEE, 86(11), 2278-2324.
  2. Krizhevsky, A., & Hinton, G. (2009). "Learning multiple layers of features from tiny images." Technical Report, University of Toronto.
  3. He, K., Zhang, X., Ren, S., & Sun, J. (2016). "Deep residual learning for image recognition." CVPR.
  4. Smith, L. N. (2017). "Cyclical learning rates for training neural networks." WACV.
  5. Selvaraju, R. R., et al. (2017). "Grad-CAM: Visual explanations from deep networks via gradient-based localization." ICCV.
  6. Ioffe, S., & Szegedy, C. (2015). "Batch normalization: Accelerating deep network training by reducing internal covariate shift." ICML.

Disclaimer