🖼️ Lab 14: Image Data

🖼️ Lab 14: Image Data#

BINF 4002 – Machine Learning for Health

Learning Objectives#

  1. Understand how images are represented as multi-dimensional arrays (height \(\times\) width \(\times\) channels)

  2. Load and explore a real medical image classification dataset

  3. Build baselines by flattening pixels and extracting hand-crafted features (HOG) + classical ML (“tabularization”)

  4. Understand convolutions — the core operation in image models — and implement one by hand

  5. Implement and train a Convolutional Neural Network (CNN) in PyTorch

  6. Apply transfer learning by fine-tuning a pre-trained ResNet on medical images

  7. Compare tabularization, CNN, and transfer-learning approaches

Why Images in Healthcare?#

Medical imaging is one of the largest and most impactful areas of health AI:

  • Radiology: chest X-rays, CT scans, MRI — detecting pneumonia, tumors, fractures

  • Pathology: microscopy of tissue biopsies — grading cancer, identifying cell types

  • Dermatology: photographs of skin lesions — classifying melanoma vs. benign nevi

  • Ophthalmology: retinal fundus images — detecting diabetic retinopathy

In this lab, we use PneumoniaMNIST, a chest X-ray dataset for binary pneumonia detection. Each image is a 28\(\times\)28 grayscale X-ray, and the task is to classify whether the patient has pneumonia. This is derived from the widely-used Chest X-Ray dataset and packaged in the MedMNIST benchmark for easy, standardized use.

Dataset: PneumoniaMNIST#

  • Source: Pediatric chest X-rays from Guangzhou Women and Children’s Medical Center

  • Size: ~5,800 images (train/val/test pre-split)

  • Labels: Binary — Normal (0) vs. Pneumonia (1)

  • Resolution: 28 \(\times\) 28 grayscale (downsampled from originals for efficiency)

Set-up#

Install Dependencies#

This lab requires PyTorch, torchvision, and MedMNIST. Run the cell below to install everything. (This may take 1-2 minutes on Colab.)

# -- Install dependencies -------------------------------------------------------
import subprocess, sys

def install(pkg):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pkg])

try:
    import torch, torchvision
    print(f"PyTorch {torch.__version__}, torchvision {torchvision.__version__}")
except ImportError:
    install("torch"); install("torchvision")
    import torch, torchvision
    print(f"PyTorch {torch.__version__}, torchvision {torchvision.__version__}")

try:
    import medmnist
    print(f"MedMNIST {medmnist.__version__}")
except ImportError:
    install("medmnist")
    import medmnist
    print(f"MedMNIST {medmnist.__version__}")

try:
    from skimage.feature import hog
    print("scikit-image OK")
except ImportError:
    install("scikit-image")
    from skimage.feature import hog

try:
    import sklearn
    print(f"scikit-learn {sklearn.__version__}")
except ImportError:
    install("scikit-learn")

Imports#

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')

plt.rcParams['figure.dpi'] = 120
plt.rcParams['axes.spines.top'] = False
plt.rcParams['axes.spines.right'] = False

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import torchvision.transforms as transforms
import torchvision.models as models

from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.metrics import roc_auc_score, roc_curve
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA

from skimage.feature import hog

import medmnist
from medmnist import PneumoniaMNIST

print("All imports successful")
print(f"  PyTorch: {torch.__version__}")
print(f"  CUDA available: {torch.cuda.is_available()}")

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

Part 1 — What Is Image Data?#

A digital image is a grid of pixel values:

  • Grayscale: a 2D array of shape \((H, W)\) — each pixel is a brightness value (0-255)

  • Color (RGB): a 3D array of shape \((H, W, 3)\) — each pixel has red, green, and blue channels

For deep learning, images are typically represented as tensors of shape \((C, H, W)\):

  • \(C\) = number of channels (1 for grayscale, 3 for RGB)

  • \(H\) = height in pixels

  • \(W\) = width in pixels

Key Properties of Image Data#

Property

Implication

Spatial structure

Nearby pixels are correlated (a nose pixel is next to a cheek pixel)

Translation invariance

A tumor in the top-left looks the same as one in the bottom-right

Fixed size (usually)

Unlike text or graphs, images are typically resized to a standard resolution

High dimensionality

A 224\(\times\)224 RGB image has 150,528 raw features

These properties motivate convolutional neural networks, which exploit spatial structure through local filters that slide across the image.

# -- Build a toy image to understand the representation -------------------------
# Create a simple 8x8 "image" with a pattern
toy = np.zeros((8, 8), dtype=np.float32)
# Draw a cross pattern
toy[3, :] = 1.0  # horizontal bar
toy[:, 3] = 1.0  # vertical bar
toy[3, 3] = 1.0  # center

fig, axes = plt.subplots(1, 3, figsize=(14, 4))

# The image
axes[0].imshow(toy, cmap='gray', vmin=0, vmax=1)
axes[0].set_title("8x8 Grayscale Image", fontsize=12)
axes[0].set_xlabel("Width (pixels)")
axes[0].set_ylabel("Height (pixels)")

# The raw pixel values
im = axes[1].imshow(toy, cmap='gray', vmin=0, vmax=1)
for i in range(8):
    for j in range(8):
        axes[1].text(j, i, f'{toy[i,j]:.0f}', ha='center', va='center',
                     fontsize=9, color='red' if toy[i,j] > 0.5 else 'gray')
axes[1].set_title("Pixel Values (0=black, 1=white)", fontsize=12)

# Flattened as a vector
flat = toy.flatten()
axes[2].bar(range(len(flat)), flat, color='#3498db', width=1.0)
axes[2].set_xlabel("Pixel Index (row-major)")
axes[2].set_ylabel("Value")
axes[2].set_title(f"Flattened: {len(flat)}-dim Vector", fontsize=12)
axes[2].set_xlim(-1, len(flat))

plt.tight_layout()
plt.show()

print(f"Image shape: {toy.shape}")
print(f"Flattened shape: {flat.shape}")
print(f"Flattening discards 2D spatial structure!")

# -- What a real medical image looks like at different resolutions ---------------
# We'll load one X-ray and show it at different scales
download = True
temp_dataset = PneumoniaMNIST(split='train', download=download, size=28)
sample_img = temp_dataset[0][0]  # PIL Image
sample_arr = np.array(sample_img)

fig, axes = plt.subplots(1, 4, figsize=(16, 4))

# Original 28x28
axes[0].imshow(sample_arr, cmap='gray')
axes[0].set_title(f"28x28 ({28*28:,} pixels)", fontsize=11)

# Downsampled to 14x14
from PIL import Image
img_14 = np.array(Image.fromarray(sample_arr.squeeze()).resize((14, 14)))
axes[1].imshow(img_14, cmap='gray')
axes[1].set_title(f"14x14 ({14*14:,} pixels)", fontsize=11)

# Downsampled to 7x7
img_7 = np.array(Image.fromarray(sample_arr.squeeze()).resize((7, 7)))
axes[2].imshow(img_7, cmap='gray')
axes[2].set_title(f"7x7 ({7*7:,} pixels)", fontsize=11)

# Show pixel grid at 7x7
axes[3].imshow(img_7, cmap='gray')
for i in range(7):
    for j in range(7):
        axes[3].text(j, i, f'{img_7[i,j]:.0f}', ha='center', va='center',
                     fontsize=7, color='yellow')
axes[3].set_title("7x7 with pixel values", fontsize=11)

for ax in axes:
    ax.set_xticks([]); ax.set_yticks([])

plt.suptitle("Resolution vs. Information: Chest X-ray at Different Scales", fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

🤔 Reflection 1.1 — Images vs. Other Data Modalities#

  1. A 28x28 grayscale image has 784 pixels. If we flatten it to a 784-dimensional vector and feed it to logistic regression, what spatial information is lost? Give a concrete example of two images that would be hard to distinguish after flattening.

  2. Compare image data to the modalities from previous labs:

    • In graphs (Lab 12), each sample had variable numbers of nodes. Images have fixed size.

    • In text (Lab 13), order matters (word sequence). Images have 2D spatial structure.

    • How do these differences affect the choice of architecture?

  3. A 224x224 RGB image has 150,528 raw features — more than most tabular datasets have samples. Why doesn’t this immediately cause overfitting when we use a CNN? (Hint: think about parameter sharing in convolutions.)

  4. Medical images often contain irrelevant regions (e.g., the black border around an X-ray, or a ruler in a dermoscopy image). How might these artifacts affect a model? How could you mitigate this?

# == SOLUTION -- Reflection 1.1 ===================================================
# 1. Flattening destroys the 2D adjacency structure. Two images with the same object
#    in different locations (e.g., a lung opacity in the upper-left vs. lower-right)
#    would have very different flat vectors even though the clinical finding is identical.
#    Also, two images with similar pixel histograms but different spatial arrangements
#    (e.g., normal lung vs. pneumonia with the same overall brightness) would look
#    similar in flattened space.
#
# 2. Images have FIXED size and REGULAR 2D grid structure -- this enables weight
#    sharing via convolutions (same filter applied at every position). Graphs have
#    VARIABLE size and IRREGULAR topology -- requiring message passing that adapts
#    to connectivity. Text has VARIABLE length and 1D SEQUENTIAL structure -- handled
#    by RNNs or self-attention. The regularity of image grids is what makes CNNs so
#    efficient: one 3x3 filter with 9 parameters can scan an entire 224x224 image.
#
# 3. CNNs exploit parameter sharing: a 3x3 convolution filter has only 9 weights but
#    is applied at every spatial position. A layer with 64 3x3 filters has only
#    64 * 9 * C = 576 * C parameters, regardless of image resolution. This is vastly
#    fewer than the ~150K weights a fully-connected layer would need. Additionally,
#    pooling layers progressively reduce spatial dimensions, and data augmentation
#    (flipping, rotating, cropping) effectively multiplies the training set.
#
# 4. Models can learn spurious correlations with artifacts. For example, if all
#    pneumonia X-rays come from Hospital A (which uses a specific scanner leaving a
#    watermark), the model might learn to detect the watermark rather than pneumonia.
#    Mitigation: crop borders, standardize preprocessing across sources, use data
#    augmentation, and validate on data from different institutions (external validation).
print("See comments above for solution.")

Part 2 — Loading and Exploring Medical Image Data#

We load PneumoniaMNIST from the MedMNIST benchmark. The dataset comes pre-split into train/validation/test sets and includes 28x28 grayscale chest X-ray images.

# -- Load PneumoniaMNIST ---------------------------------------------------------
# MedMNIST provides a standardized interface similar to torchvision datasets

train_dataset = PneumoniaMNIST(split='train', download=True, size=28)
val_dataset = PneumoniaMNIST(split='val', download=True, size=28)
test_dataset = PneumoniaMNIST(split='test', download=True, size=28)

# Extract arrays for easier manipulation
X_train_img = np.array([np.array(train_dataset[i][0]).squeeze() for i in range(len(train_dataset))])
y_train_raw = np.array([train_dataset[i][1].item() for i in range(len(train_dataset))])

X_val_img = np.array([np.array(val_dataset[i][0]).squeeze() for i in range(len(val_dataset))])
y_val_raw = np.array([val_dataset[i][1].item() for i in range(len(val_dataset))])

X_test_img = np.array([np.array(test_dataset[i][0]).squeeze() for i in range(len(test_dataset))])
y_test_raw = np.array([test_dataset[i][1].item() for i in range(len(test_dataset))])

# Binarize labels (PneumoniaMNIST: 0=normal, 1=pneumonia)
y_train = (y_train_raw > 0).astype(int)
y_val = (y_val_raw > 0).astype(int)
y_test = (y_test_raw > 0).astype(int)

print(f"Train: {X_train_img.shape} images, {y_train.sum()}/{len(y_train)} pneumonia ({y_train.mean():.1%})")
print(f"Val:   {X_val_img.shape} images, {y_val.sum()}/{len(y_val)} pneumonia ({y_val.mean():.1%})")
print(f"Test:  {X_test_img.shape} images, {y_test.sum()}/{len(y_test)} pneumonia ({y_test.mean():.1%})")
print(f"\nPixel value range: [{X_train_img.min()}, {X_train_img.max()}]")
print(f"Image dtype: {X_train_img.dtype}")

# -- Visualize sample images from each class -------------------------------------
fig, axes = plt.subplots(2, 8, figsize=(16, 5))

# Normal examples
normal_idx = np.where(y_train == 0)[0]
for i in range(8):
    axes[0, i].imshow(X_train_img[normal_idx[i]], cmap='gray')
    axes[0, i].set_xticks([]); axes[0, i].set_yticks([])
    if i == 0:
        axes[0, i].set_ylabel("Normal", fontsize=12, fontweight='bold')

# Pneumonia examples
pneumonia_idx = np.where(y_train == 1)[0]
for i in range(8):
    axes[1, i].imshow(X_train_img[pneumonia_idx[i]], cmap='gray')
    axes[1, i].set_xticks([]); axes[1, i].set_yticks([])
    if i == 0:
        axes[1, i].set_ylabel("Pneumonia", fontsize=12, fontweight='bold')

plt.suptitle("Sample Chest X-rays: Normal vs. Pneumonia", fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

# -- Dataset statistics -----------------------------------------------------------
fig, axes = plt.subplots(1, 3, figsize=(15, 4))

# Label distribution
for idx, (name, y) in enumerate([('Train', y_train), ('Val', y_val), ('Test', y_test)]):
    counts = [np.sum(y == 0), np.sum(y == 1)]
    axes[0].bar([idx - 0.2, idx + 0.2], counts, width=0.35,
               color=['#3498db', '#e74c3c'])
    for j, v in enumerate(counts):
        axes[0].text(idx - 0.2 + j * 0.4, v + 10, str(v), ha='center', fontsize=9)
axes[0].set_xticks([0, 1, 2]); axes[0].set_xticklabels(['Train', 'Val', 'Test'])
axes[0].set_ylabel('Count')
axes[0].set_title('Label Distribution')
axes[0].legend(['Normal', 'Pneumonia'], loc='upper right')

# Pixel intensity distribution by class
normal_pixels = X_train_img[y_train == 0].flatten()
pneumonia_pixels = X_train_img[y_train == 1].flatten()
axes[1].hist(normal_pixels, bins=50, alpha=0.5, color='#3498db',
             label='Normal', density=True)
axes[1].hist(pneumonia_pixels, bins=50, alpha=0.5, color='#e74c3c',
             label='Pneumonia', density=True)
axes[1].set_xlabel('Pixel Value')
axes[1].set_ylabel('Density')
axes[1].set_title('Pixel Intensity Distribution')
axes[1].legend()

# Mean images per class
mean_normal = X_train_img[y_train == 0].mean(axis=0)
mean_pneumonia = X_train_img[y_train == 1].mean(axis=0)
diff = mean_pneumonia.astype(float) - mean_normal.astype(float)

# Show difference image
im = axes[2].imshow(diff, cmap='RdBu_r', vmin=-30, vmax=30)
axes[2].set_title('Mean Difference\n(Pneumonia - Normal)', fontsize=11)
axes[2].set_xticks([]); axes[2].set_yticks([])
plt.colorbar(im, ax=axes[2], shrink=0.8)

plt.tight_layout()
plt.show()

print(f"\nPixel statistics:")
print(f"  Normal mean brightness: {X_train_img[y_train == 0].mean():.1f}")
print(f"  Pneumonia mean brightness: {X_train_img[y_train == 1].mean():.1f}")

🤔 Reflection 2.1 — Understanding Medical Image Data#

  1. Look at the mean difference image (Pneumonia minus Normal). Which regions of the chest X-ray show the largest differences? Does this match your clinical intuition about where pneumonia manifests?

  2. The dataset is imbalanced (~75% pneumonia). In clinical practice, pneumonia prevalence varies widely (5% in outpatient clinics, 50% in ICUs). How would you adapt your model evaluation to account for different deployment settings?

  3. At 28x28 resolution, fine details (like small nodules or subtle infiltrates) are lost. What is the trade-off between resolution and computational cost? At what point might higher resolution stop helping?

  4. The pixel intensity distributions overlap substantially between classes. What does this tell you about the difficulty of classification from raw pixel values alone?

# == SOLUTION -- Reflection 2.1 ===================================================
# 1. The largest differences are typically in the lung fields (central and lower
#    regions of the image). Pneumonia causes opacification (whitening) of the affected
#    lung areas, so pneumonia X-rays tend to be brighter in the lung regions compared
#    to normal X-rays which show clear, dark lung fields. This matches clinical
#    knowledge that pneumonia appears as consolidation or infiltrates in the lungs.
#
# 2. The key insight is that model performance metrics depend on prevalence.
#    AUROC is prevalence-independent, so it's useful for comparing models. But
#    clinical deployment metrics (PPV, NPV) depend strongly on prevalence.
#    A model with 90% sensitivity and 90% specificity has PPV=75% at 25% prevalence
#    but only PPV=32% at 5% prevalence. We should report AUROC for model comparison
#    but use decision curve analysis and calibration for deployment decisions.
#
# 3. Higher resolution captures finer details but costs O(H*W) more computation per
#    layer and requires more GPU memory. For 28x28 -> 224x224, that's a 64x increase
#    in pixels. Diminishing returns set in when the additional detail doesn't improve
#    the task: for detecting large consolidations (pneumonia), 112x112 may suffice;
#    for detecting tiny nodules, 512x512+ may be needed. In practice, 224x224 is the
#    standard compromise for most classification tasks.
#
# 4. Heavy overlap in pixel distributions means global statistics (mean brightness,
#    histogram features) have limited discriminative power. The signal lies in the
#    SPATIAL PATTERN of pixel values -- not their marginal distribution. This is
#    precisely why spatial models (CNNs) outperform models that discard spatial
#    structure. The difference between pneumonia and normal is WHERE the bright
#    pixels are, not just HOW MANY bright pixels there are.
print("See comments above for solution.")

Part 3 — Baseline: Image Tabularization#

As with graphs and text, we start by converting images to fixed-length feature vectors.

Approach A — Flatten Raw Pixels#

The simplest approach: reshape each 28\(\times\)28 image into a 784-dimensional vector. This discards all spatial structure but gives us a direct tabular representation.

Approach B — Histogram of Oriented Gradients (HOG)#

HOG captures local texture and edge patterns:

  1. Compute image gradients (edge directions) at each pixel

  2. Divide the image into small cells (e.g., 7\(\times\)7 pixels)

  3. Build a histogram of gradient orientations within each cell

  4. Normalize across blocks of cells for illumination robustness

HOG preserves some spatial information (which cell contains which edges) while producing a fixed-length vector. It’s the classic hand-crafted feature for image recognition.

# This is an exercise cell with TODOs (??? placeholders) for student work.
# It is skipped in the book build to keep downstream cells runnable.
# Open the notebook in Colab to complete the exercise:
# Exercise cell skipped in book build
print("[exercise cell skipped — open in Colab to complete the TODOs]")

# == SOLUTION -- compute_hog_manual ================================================
def compute_hog_manual(image, cell_size=7, n_bins=9):
    H, W = image.shape
    img = image.astype(np.float64)

    gy, gx = np.gradient(img)
    magnitude = np.sqrt(gx**2 + gy**2)
    orientation = np.degrees(np.arctan2(gy, gx)) % 180

    n_cells_y = H // cell_size
    n_cells_x = W // cell_size
    histograms = []

    for cy in range(n_cells_y):
        for cx in range(n_cells_x):
            y_start = cy * cell_size
            x_start = cx * cell_size
            cell_mag = magnitude[y_start:y_start+cell_size, x_start:x_start+cell_size]
            cell_ori = orientation[y_start:y_start+cell_size, x_start:x_start+cell_size]

            hist, _ = np.histogram(cell_ori.flatten(), bins=n_bins, range=(0, 180),
                                   weights=cell_mag.flatten())
            histograms.append(hist)

    feature_vector = np.concatenate(histograms)
    norm = np.linalg.norm(feature_vector) + 1e-6
    feature_vector = feature_vector / norm
    return feature_vector

test_hog = compute_hog_manual(X_train_img[0].astype(np.float64))
print(f"HOG feature vector length: {len(test_hog)}")
print(f"First 10 values: {test_hog[:10].round(4)}")

# -- Visualize HOG features -------------------------------------------------------
from skimage.feature import hog as skimage_hog

fig, axes = plt.subplots(2, 4, figsize=(16, 8))

for i, idx in enumerate([normal_idx[0], normal_idx[1], pneumonia_idx[0], pneumonia_idx[1]]):
    img = X_train_img[idx]
    hog_feat, hog_img = skimage_hog(img, orientations=9, pixels_per_cell=(7, 7),
                                     cells_per_block=(1, 1), visualize=True)
    label = "Normal" if y_train[idx] == 0 else "Pneumonia"
    axes[0, i].imshow(img, cmap='gray')
    axes[0, i].set_title(f"{label}", fontsize=11)
    axes[0, i].set_xticks([]); axes[0, i].set_yticks([])
    axes[1, i].imshow(hog_img, cmap='gray')
    axes[1, i].set_title("HOG features", fontsize=11)
    axes[1, i].set_xticks([]); axes[1, i].set_yticks([])

axes[0, 0].set_ylabel("Original", fontsize=12)
axes[1, 0].set_ylabel("HOG", fontsize=12)
plt.suptitle("HOG Captures Edge Patterns in Chest X-rays", fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

# -- Build tabular feature matrices ------------------------------------------------
print("Extracting features...")

# Approach A: Flattened pixels (normalize to [0, 1])
X_train_flat = X_train_img.reshape(len(X_train_img), -1).astype(np.float32) / 255.0
X_val_flat = X_val_img.reshape(len(X_val_img), -1).astype(np.float32) / 255.0
X_test_flat = X_test_img.reshape(len(X_test_img), -1).astype(np.float32) / 255.0
print(f"  Flat pixels: {X_train_flat.shape}")

# Approach B: HOG features (using scikit-image for speed)
def extract_hog_features(images):
    features = []
    for img in images:
        feat = skimage_hog(img, orientations=9, pixels_per_cell=(7, 7),
                           cells_per_block=(2, 2), feature_vector=True)
        features.append(feat)
    return np.array(features)

X_train_hog = extract_hog_features(X_train_img)
X_val_hog = extract_hog_features(X_val_img)
X_test_hog = extract_hog_features(X_test_img)
print(f"  HOG features: {X_train_hog.shape}")

# Standardize
scaler_flat = StandardScaler()
X_train_flat_s = scaler_flat.fit_transform(X_train_flat)
X_val_flat_s = scaler_flat.transform(X_val_flat)
X_test_flat_s = scaler_flat.transform(X_test_flat)

scaler_hog = StandardScaler()
X_train_hog_s = scaler_hog.fit_transform(X_train_hog)
X_val_hog_s = scaler_hog.transform(X_val_hog)
X_test_hog_s = scaler_hog.transform(X_test_hog)

# -- Classical ML baselines --------------------------------------------------------
print("=== Tabularization Baselines ===\n")

print("-- Flattened Pixels (784-dim) --")
flat_models = {
    'LR + Pixels': LogisticRegression(max_iter=1000, C=1.0, random_state=42),
    'RF + Pixels': RandomForestClassifier(n_estimators=200, max_depth=15, random_state=42),
}
flat_results = {}
for name, model in flat_models.items():
    model.fit(X_train_flat_s, y_train)
    val_proba = model.predict_proba(X_val_flat_s)[:, 1]
    val_auc = roc_auc_score(y_val, val_proba)
    flat_results[name] = {'model': model, 'val_auc': val_auc}
    print(f"  {name:25s}  Val AUROC = {val_auc:.4f}")

print("\n-- HOG Features --")
hog_models = {
    'LR + HOG': LogisticRegression(max_iter=1000, C=1.0, random_state=42),
    'RF + HOG': RandomForestClassifier(n_estimators=200, max_depth=15, random_state=42),
    'GBT + HOG': GradientBoostingClassifier(n_estimators=200, max_depth=5, random_state=42),
}
hog_results = {}
for name, model in hog_models.items():
    model.fit(X_train_hog_s, y_train)
    val_proba = model.predict_proba(X_val_hog_s)[:, 1]
    val_auc = roc_auc_score(y_val, val_proba)
    hog_results[name] = {'model': model, 'val_auc': val_auc}
    print(f"  {name:25s}  Val AUROC = {val_auc:.4f}")

🤔 Reflection 3.1 — Tabularization Trade-offs#

  1. Compare flattened pixels vs. HOG features. Which performs better? Why does HOG capture information that raw pixels miss?

  2. The flattened pixel vector has 784 dimensions. HOG has far fewer features but (typically) performs comparably or better. What does this tell you about the information density of each representation?

  3. In Labs 6-7, we saw a similar pattern: hand-crafted features (fingerprints, TF-IDF) vs. learned features (GNNs, LSTMs). What is the common principle?

  4. For a 224\(\times\)224 RGB image, the flattened vector would have 150,528 features. Why would logistic regression struggle with this input? (Hint: think about both computational cost and the curse of dimensionality.)

# == SOLUTION -- Reflection 3.1 ===================================================
# 1. HOG typically outperforms raw pixels because it extracts MEANINGFUL patterns
#    (edges, gradients, textures) rather than raw brightness values. A pneumonia
#    infiltrate appears as a characteristic pattern of edges and intensity gradients
#    that HOG captures directly. Raw pixels are sensitive to irrelevant variations
#    (overall brightness, exact positioning) that don't relate to the diagnosis.
#
# 2. HOG features are much more information-dense: each HOG dimension encodes a
#    meaningful edge pattern in a specific region. In raw pixels, most individual
#    pixel values are uninformative (a single bright pixel tells you almost nothing).
#    This is a fundamental principle: GOOD FEATURES compress the data into dimensions
#    that are relevant to the task, discarding irrelevant variation.
#
# 3. The common principle: domain-informed feature extraction captures most of the
#    signal, often matching or approaching learned features on small datasets. Morgan
#    fingerprints capture chemical substructures, TF-IDF captures word importance,
#    HOG captures edge patterns. The advantage of learned features (GNNs, LSTMs,
#    CNNs) is that they can discover patterns that humans didn't design features for,
#    and they improve with more data.
#
# 4. With 150K features and maybe 5K training samples, logistic regression is
#    massively over-parameterized (30 parameters per sample). This leads to severe
#    overfitting. Additionally, storing and inverting a 150K x 150K matrix is
#    computationally expensive. PCA or other dimensionality reduction would help,
#    but the fundamental problem is that linear models can't exploit the spatial
#    structure that makes high-resolution images useful.
print("See comments above for solution.")

Part 4 — Understanding Convolutions#

The convolution operation is the core building block of CNNs. A small filter (kernel) slides across the image, computing a weighted sum at each position:

\[(I * K)[i, j] = \sum_{m} \sum_{n} I[i+m,\; j+n] \cdot K[m, n]\]

where \(I\) is the input image and \(K\) is the kernel (e.g., 3\(\times\)3).

Why Convolutions Work for Images#

  1. Local connectivity: each output depends only on a small local region

  2. Parameter sharing: the same kernel is used at every position → far fewer parameters

  3. Translation equivariance: the same pattern is detected regardless of its position

A CNN stacks multiple convolutional layers, progressively building from low-level features (edges, textures) to high-level features (anatomical structures, pathological patterns).

# This is an exercise cell with TODOs (??? placeholders) for student work.
# It is skipped in the book build to keep downstream cells runnable.
# Open the notebook in Colab to complete the exercise:
# Exercise cell skipped in book build
print("[exercise cell skipped — open in Colab to complete the TODOs]")

# == SOLUTION -- conv2d_manual =====================================================
def conv2d_manual(image, kernel):
    H, W = image.shape
    kH, kW = kernel.shape
    out_H = H - kH + 1
    out_W = W - kW + 1
    output = np.zeros((out_H, out_W))

    for i in range(out_H):
        for j in range(out_W):
            region = image[i:i+kH, j:j+kW]
            output[i, j] = np.sum(region * kernel)

    return output

edge_kernel = np.array([[-1, -1, -1],
                         [ 0,  0,  0],
                         [ 1,  1,  1]], dtype=np.float64)

test_img = X_train_img[0].astype(np.float64)
edge_map = conv2d_manual(test_img, edge_kernel)
print(f"Input shape:  {test_img.shape}")
print(f"Kernel shape: {edge_kernel.shape}")
print(f"Output shape: {edge_map.shape}")

# -- Visualize different kernels and their effects ---------------------------------
kernels = {
    'Horizontal
Edges': np.array([[-1,-1,-1],[0,0,0],[1,1,1]], dtype=np.float64),
    'Vertical
Edges': np.array([[-1,0,1],[-1,0,1],[-1,0,1]], dtype=np.float64),
    'Sharpen': np.array([[0,-1,0],[-1,5,-1],[0,-1,0]], dtype=np.float64),
    'Blur': np.ones((3,3), dtype=np.float64) / 9.0,
}

fig, axes = plt.subplots(2, len(kernels) + 1, figsize=(16, 7))

# Show original image
for row in range(2):
    axes[row, 0].imshow(X_train_img[pneumonia_idx[row]], cmap='gray')
    label = "Pneumonia" if y_train[pneumonia_idx[row]] else "Normal"
    axes[row, 0].set_title(f"Original\n({label})" if row == 0 else "Original", fontsize=10)
    axes[row, 0].set_xticks([]); axes[row, 0].set_yticks([])

# Apply each kernel
for k_idx, (name, kernel) in enumerate(kernels.items()):
    for row in range(2):
        img = X_train_img[pneumonia_idx[row]].astype(np.float64)
        result = conv2d_manual(img, kernel)
        axes[row, k_idx + 1].imshow(result, cmap='gray')
        if row == 0:
            axes[row, k_idx + 1].set_title(name, fontsize=10)
        axes[row, k_idx + 1].set_xticks([]); axes[row, k_idx + 1].set_yticks([])

plt.suptitle("Convolving X-rays with Different Kernels", fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

# Show the kernels themselves
fig, axes = plt.subplots(1, len(kernels), figsize=(12, 3))
for k_idx, (name, kernel) in enumerate(kernels.items()):
    im = axes[k_idx].imshow(kernel, cmap='RdBu_r', vmin=-2, vmax=2)
    axes[k_idx].set_title(name.replace('\n', ' '), fontsize=10)
    for i in range(3):
        for j in range(3):
            axes[k_idx].text(j, i, f'{kernel[i,j]:.1f}', ha='center', va='center', fontsize=9)
plt.suptitle("The Kernels (Filters)", fontsize=12)
plt.tight_layout()
plt.show()

🤔 Reflection 4.1 — Convolutions and Spatial Features#

  1. The horizontal edge detector kernel is [[-1,-1,-1],[0,0,0],[1,1,1]]. Explain intuitively why this detects horizontal edges. What would happen at a pixel where the top neighbors are dark and the bottom neighbors are bright?

  2. A single 3\(\times\)3 kernel has 9 learnable parameters. A fully connected layer connecting a 28\(\times\)28 image to 28\(\times\)28 output has \(784^2 = 614{,}656\) parameters. Calculate the parameter reduction factor of convolution. How does this relate to the concept of inductive bias?

  3. After convolution with a 3\(\times\)3 kernel (no padding), the output is 26\(\times\)26. After another 3\(\times\)3 convolution, it becomes 24\(\times\)24. What is the effective receptive field of a pixel in the output of two consecutive 3\(\times\)3 convolutions? How does this relate to message passing in GNNs?

  4. Looking at the convolution outputs, which kernel(s) seem most useful for detecting pneumonia patterns? What kinds of patterns would a CNN learn to detect in its first few layers vs. deeper layers?

# == SOLUTION -- Reflection 4.1 ===================================================
# 1. The kernel computes (bottom row sum) - (top row sum) at each position.
#    At a horizontal edge where top is dark (low values) and bottom is bright
#    (high values): output = (high+high+high) - (low+low+low) = large positive.
#    At a horizontal edge with reverse polarity: large negative. In uniform
#    regions: top ~= bottom, so output ~= 0. This is why edges "light up."
#
# 2. Reduction factor: 614,656 / 9 = 68,295x fewer parameters. This massive
#    reduction is the core of CNN's inductive bias: we assume that (a) the same
#    local pattern is relevant everywhere in the image (translation equivariance),
#    and (b) only local relationships matter at each layer. If this assumption
#    matches the data (it does for images!), we get much better generalization
#    with far fewer parameters. If it doesn't match (e.g., for tabular data where
#    "position" is meaningless), convolutions hurt.
#
# 3. Two consecutive 3x3 convolutions give an effective receptive field of 5x5:
#    the first layer sees 3x3, and the second layer sees 3x3 of those, covering
#    5x5 of the original input. This is exactly analogous to message passing in
#    GNNs: one layer = 1-hop neighborhood, two layers = 2-hop neighborhood. Three
#    3x3 layers have a 7x7 receptive field. Deep CNNs have very large receptive
#    fields, and the same over-smoothing concerns from GNNs can apply (though
#    pooling and striding mitigate this).
#
# 4. Edge detectors (horizontal, vertical) are most directly useful -- they outline
#    anatomical structures (ribs, diaphragm, heart border) and abnormalities.
#    First layers learn simple edges and textures. Middle layers combine edges into
#    parts (lung boundaries, rib patterns). Deeper layers detect semantic concepts
#    (consolidation, pleural effusion, cardiomegaly). This hierarchical feature
#    learning is the key advantage of deep CNNs over HOG.
print("See comments above for solution.")

Part 5 — Training a Convolutional Neural Network#

Now we build and train a CNN from scratch in PyTorch. A typical CNN for classification has:

  1. Convolutional blocks: Conv2d \(\to\) ReLU \(\to\) MaxPool (repeated 2-3 times)

  2. Flatten: reshape the 2D feature maps to a 1D vector

  3. Classifier head: Linear \(\to\) ReLU \(\to\) Linear \(\to\) Sigmoid

Our Architecture#

Input (1, 28, 28)
  --> Conv2d(1, 32, 3, padding=1) --> ReLU --> MaxPool2d(2)   -> (32, 14, 14)
  --> Conv2d(32, 64, 3, padding=1) --> ReLU --> MaxPool2d(2)  -> (64, 7, 7)
  --> Conv2d(64, 128, 3, padding=1) --> ReLU --> MaxPool2d(2) -> (128, 3, 3)
  --> Flatten                                                 -> (1152,)
  --> Linear(1152, 128) --> ReLU --> Dropout
  --> Linear(128, 1) --> Sigmoid

# This is an exercise cell with TODOs (??? placeholders) for student work.
# It is skipped in the book build to keep downstream cells runnable.
# Open the notebook in Colab to complete the exercise:
# Exercise cell skipped in book build
print("[exercise cell skipped — open in Colab to complete the TODOs]")

# == SOLUTION -- SimpleCNN =========================================================
class SimpleCNN(nn.Module):
    def __init__(self, dropout=0.3):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 3 * 3, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 1),
        )

    def forward(self, x):
        features = self.features(x)
        out = torch.sigmoid(self.classifier(features)).squeeze(-1)
        return out

cnn_model = SimpleCNN()
print(cnn_model)
dummy = torch.randn(2, 1, 28, 28)
dummy_out = cnn_model(dummy)
print(f"\nInput shape:  {dummy.shape}")
print(f"Output shape: {dummy_out.shape}")
print(f"Total parameters: {sum(p.numel() for p in cnn_model.parameters()):,}")

# -- Prepare PyTorch dataloaders ----------------------------------------------------
# Normalize images to [0, 1] and add channel dimension

def prepare_image_tensors(images, labels):
    X = torch.tensor(images, dtype=torch.float32).unsqueeze(1) / 255.0  # (N, 1, 28, 28)
    y = torch.tensor(labels, dtype=torch.float32)
    return TensorDataset(X, y)

train_ds = prepare_image_tensors(X_train_img, y_train)
val_ds = prepare_image_tensors(X_val_img, y_val)
test_ds = prepare_image_tensors(X_test_img, y_test)

train_dl = DataLoader(train_ds, batch_size=64, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=64, shuffle=False)
test_dl = DataLoader(test_ds, batch_size=64, shuffle=False)

print(f"Train batches: {len(train_dl)}, Val: {len(val_dl)}, Test: {len(test_dl)}")
batch_x, batch_y = next(iter(train_dl))
print(f"Batch X shape: {batch_x.shape}, Batch y shape: {batch_y.shape}")

# -- Training and evaluation functions ---------------------------------------------
def train_epoch_cnn(model, loader, optimizer, criterion):
    model.train()
    total_loss, n = 0, 0
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()
        out = model(X)
        loss = criterion(out, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * len(y)
        n += len(y)
    return total_loss / n

@torch.no_grad()
def evaluate_cnn(model, loader):
    model.eval()
    all_preds, all_labels = [], []
    total_loss, n = 0, 0
    criterion = nn.BCELoss()
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        out = model(X)
        total_loss += criterion(out, y).item() * len(y)
        all_preds.append(out.cpu().numpy())
        all_labels.append(y.cpu().numpy())
        n += len(y)
    preds = np.concatenate(all_preds)
    labels = np.concatenate(all_labels)
    auc = roc_auc_score(labels, preds)
    return auc, total_loss / n, preds, labels

def train_cnn_model(model, train_dl, val_dl, lr=1e-3, epochs=25, patience=7):
    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.BCELoss()

    history = {'train_loss': [], 'val_loss': [], 'train_auc': [], 'val_auc': []}
    best_val_auc, best_state, wait = 0, None, 0

    for epoch in range(1, epochs + 1):
        train_loss = train_epoch_cnn(model, train_dl, optimizer, criterion)
        train_auc, _, _, _ = evaluate_cnn(model, train_dl)
        val_auc, val_loss, _, _ = evaluate_cnn(model, val_dl)

        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['train_auc'].append(train_auc)
        history['val_auc'].append(val_auc)

        if val_auc > best_val_auc:
            best_val_auc = val_auc
            best_state = {k: v.clone() for k, v in model.state_dict().items()}
            wait = 0
        else:
            wait += 1

        if epoch % 5 == 0 or epoch == 1:
            print(f"Epoch {epoch:2d} | Train Loss: {train_loss:.4f} | "
                  f"Train AUC: {train_auc:.4f} | Val AUC: {val_auc:.4f} | Best: {best_val_auc:.4f}")

        if wait >= patience:
            print(f"Early stopping at epoch {epoch}")
            break

    model.load_state_dict(best_state)
    return model, history

# -- Train the CNN ----------------------------------------------------------------
print("=" * 60)
print("Training SimpleCNN from scratch")
print("=" * 60)

cnn_model = SimpleCNN(dropout=0.3)
cnn_model, cnn_history = train_cnn_model(cnn_model, train_dl, val_dl,
                                          lr=1e-3, epochs=25, patience=7)

# Plot training curves
fig, axes = plt.subplots(1, 2, figsize=(13, 4))

axes[0].plot(cnn_history['train_loss'], label='Train', color='#3498db')
axes[0].plot(cnn_history['val_loss'], label='Val', color='#e74c3c')
axes[0].set_xlabel('Epoch'); axes[0].set_ylabel('BCE Loss')
axes[0].set_title('CNN -- Loss Curves'); axes[0].legend()

axes[1].plot(cnn_history['train_auc'], label='Train', color='#3498db')
axes[1].plot(cnn_history['val_auc'], label='Val', color='#e74c3c')
axes[1].set_xlabel('Epoch'); axes[1].set_ylabel('AUROC')
axes[1].set_title('CNN -- AUROC Curves'); axes[1].legend()

plt.tight_layout()
plt.show()

# -- Visualize learned convolutional filters ----------------------------------------
# Show the 32 filters from the first convolutional layer

first_conv = list(cnn_model.features.children())[0]
filters = first_conv.weight.data.cpu().numpy()  # shape: (32, 1, 3, 3)

fig, axes = plt.subplots(4, 8, figsize=(14, 7))
for i in range(32):
    ax = axes[i // 8, i % 8]
    f = filters[i, 0]
    ax.imshow(f, cmap='RdBu_r', vmin=-0.5, vmax=0.5)
    ax.set_xticks([]); ax.set_yticks([])
    ax.set_title(f'F{i}', fontsize=8)

plt.suptitle("Learned First-Layer Filters (3x3 kernels)", fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

print("Notice: some filters look like edge detectors (similar to our hand-crafted kernels),")
print("while others capture patterns we might not have thought to design manually.")

🤔 Reflection 5.1 — CNN Architecture and Training#

  1. The CNN has ~170K parameters. Compare this to logistic regression on flattened pixels (784 parameters). The CNN has ~200x more parameters but typically achieves much better performance. Why doesn’t it overfit badly? (Hint: think about the inductive bias of convolutions and the role of pooling.)

  2. Look at the learned first-layer filters. Some resemble edge detectors (like our hand-crafted kernels from Part 4). Why is this? What does this tell you about the relationship between hand-crafted and learned features?

  3. MaxPool2d(2) halves the spatial resolution at each layer: 28\(\to\)14\(\to\)7\(\to\)3. What information is lost through pooling? Why is this loss acceptable for classification but might be problematic for other tasks (like segmentation)?

  4. We used padding=1 in each convolution to maintain spatial dimensions before pooling. What would happen without padding? How many pixels would we lose per layer?

# == SOLUTION -- Reflection 5.1 ===================================================
# 1. The CNN's inductive bias matches image structure: (a) parameter SHARING means
#    each 3x3 filter has only 9 weights applied at every position -- the model doesn't
#    learn separate weights for "tumor in top-left" vs "tumor in bottom-right."
#    (b) LOCALITY means each layer only combines nearby pixels, preventing the model
#    from learning arbitrary long-range correlations that might be noise. (c) POOLING
#    progressively reduces spatial dimensions, forcing the model to build abstract
#    representations. Together, these constraints dramatically reduce effective model
#    complexity despite the high parameter count.
#
# 2. Learning converges to edge-like filters because edges are the OPTIMAL first-level
#    features for most visual tasks -- they're the building blocks from which all higher-
#    level patterns are composed. This is a beautiful validation of the hand-crafted
#    features: decades of computer vision research arrived at the same features that
#    gradient descent discovers automatically. The advantage of learning is that the
#    CNN also discovers filters we WOULDN'T have designed (subtle texture patterns,
#    frequency-specific detectors).
#
# 3. MaxPool keeps the maximum activation in each 2x2 region, discarding the exact
#    position within that region. This provides translation INVARIANCE (slightly shifted
#    features produce the same output) but loses precise spatial information. For
#    classification ("is there pneumonia?"), we only need to know IF a pattern exists,
#    not exactly WHERE. For segmentation ("which pixels are pneumonia?"), we need to
#    preserve spatial resolution -- which is why architectures like U-Net use skip
#    connections to recover spatial detail.
#
# 4. Without padding, each 3x3 convolution shrinks the spatial dimension by 2 (one
#    pixel on each side). Three unpooled 3x3 convolutions: 28 -> 26 -> 24 -> 22. With
#    pooling interleaved: 28 -> 26 -> 13 -> 11 -> 5 -> 3 -> 1. This aggressive shrinking
#    limits the number of layers we can stack and loses border information. Padding
#    preserves the spatial dimensions, giving pooling more resolution to work with.
print("See comments above for solution.")

Part 6 — Transfer Learning with Pre-trained ResNet#

Just as DistilBERT (Lab 7) was pre-trained on large text corpora, CNNs can be pre-trained on large image datasets. ImageNet (1.2M natural images, 1000 classes) is the standard pre-training dataset for vision.

The Transfer Learning Pipeline#

  1. Load a ResNet-18 pre-trained on ImageNet (11.7M parameters)

  2. Replace the final classification layer (1000 classes \(\to\) 1 class)

  3. Adapt the input (our 1-channel grayscale \(\to\) 3-channel by replication)

  4. Fine-tune with a small learning rate to preserve pre-trained features

Why Does This Work?#

Pre-trained ImageNet features transfer surprisingly well to medical images:

  • First layers detect edges and textures — universal across all image types

  • Middle layers detect patterns and parts — often transferable

  • Last layers detect ImageNet-specific objects — replaced with our task head

Clinical parallel: Models like CheXNet (DenseNet pre-trained on ImageNet, fine-tuned on chest X-rays) achieve radiologist-level pneumonia detection.

# -- Fine-tune a pre-trained ResNet-18 ---------------------------------------------
print("Loading pre-trained ResNet-18...")

# Load pre-trained model
resnet = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

# Modify first conv layer to accept 1-channel input
# Strategy: average the 3-channel weights into 1 channel
original_conv = resnet.conv1
resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
with torch.no_grad():
    resnet.conv1.weight = nn.Parameter(original_conv.weight.mean(dim=1, keepdim=True))

# Replace final classification layer
resnet.fc = nn.Sequential(
    nn.Linear(resnet.fc.in_features, 1),
    nn.Sigmoid()
)

print(f"ResNet-18 parameters: {sum(p.numel() for p in resnet.parameters()):,}")

# Freeze early layers (only fine-tune the last block + FC)
for name, param in resnet.named_parameters():
    if 'layer4' not in name and 'fc' not in name:
        param.requires_grad = False

trainable = sum(p.numel() for p in resnet.parameters() if p.requires_grad)
frozen = sum(p.numel() for p in resnet.parameters() if not p.requires_grad)
print(f"Trainable: {trainable:,} | Frozen: {frozen:,}")

# -- Prepare data for ResNet (needs upsampling to 224x224) -------------------------
# ResNet expects 224x224; we upsample our 28x28 images

def prepare_resnet_tensors(images, labels):
    X = torch.tensor(images, dtype=torch.float32).unsqueeze(1) / 255.0  # (N, 1, 28, 28)
    # Upsample to 224x224 using bilinear interpolation
    X = F.interpolate(X, size=(224, 224), mode='bilinear', align_corners=False)
    # Normalize with ImageNet statistics (adapted for single-channel)
    X = (X - 0.449) / 0.226  # approximate ImageNet mean/std for grayscale
    y = torch.tensor(labels, dtype=torch.float32)
    return TensorDataset(X, y)

print("Preparing upsampled datasets (this may take a moment)...")
train_ds_resnet = prepare_resnet_tensors(X_train_img, y_train)
val_ds_resnet = prepare_resnet_tensors(X_val_img, y_val)
test_ds_resnet = prepare_resnet_tensors(X_test_img, y_test)

train_dl_resnet = DataLoader(train_ds_resnet, batch_size=32, shuffle=True)
val_dl_resnet = DataLoader(val_ds_resnet, batch_size=32, shuffle=False)
test_dl_resnet = DataLoader(test_ds_resnet, batch_size=32, shuffle=False)

print(f"Resnet batch shape: {next(iter(train_dl_resnet))[0].shape}")

# -- Train ResNet ------------------------------------------------------------------
# Use the same training infrastructure but with smaller learning rate

def forward_resnet(model, X):
    return model(X).squeeze(-1)

def train_epoch_resnet(model, loader, optimizer, criterion):
    model.train()
    total_loss, n = 0, 0
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()
        out = forward_resnet(model, X)
        loss = criterion(out, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * len(y)
        n += len(y)
    return total_loss / n

@torch.no_grad()
def evaluate_resnet(model, loader):
    model.eval()
    all_preds, all_labels = [], []
    total_loss, n = 0, 0
    criterion = nn.BCELoss()
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        out = forward_resnet(model, X)
        total_loss += criterion(out, y).item() * len(y)
        all_preds.append(out.cpu().numpy())
        all_labels.append(y.cpu().numpy())
        n += len(y)
    preds = np.concatenate(all_preds)
    labels = np.concatenate(all_labels)
    auc = roc_auc_score(labels, preds)
    return auc, total_loss / n, preds, labels

print("=" * 60)
print("Fine-tuning ResNet-18 (pre-trained on ImageNet)")
print("=" * 60)

resnet = resnet.to(device)
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, resnet.parameters()),
                             lr=1e-4, weight_decay=1e-4)
criterion = nn.BCELoss()

resnet_history = {'train_loss': [], 'val_loss': [], 'train_auc': [], 'val_auc': []}
best_val_auc, best_state, wait = 0, None, 0
PATIENCE = 5

for epoch in range(1, 16):
    train_loss = train_epoch_resnet(resnet, train_dl_resnet, optimizer, criterion)
    train_auc, _, _, _ = evaluate_resnet(resnet, train_dl_resnet)
    val_auc, val_loss, _, _ = evaluate_resnet(resnet, val_dl_resnet)

    resnet_history['train_loss'].append(train_loss)
    resnet_history['val_loss'].append(val_loss)
    resnet_history['train_auc'].append(train_auc)
    resnet_history['val_auc'].append(val_auc)

    if val_auc > best_val_auc:
        best_val_auc = val_auc
        best_state = {k: v.clone() for k, v in resnet.state_dict().items()}
        wait = 0
    else:
        wait += 1

    print(f"Epoch {epoch:2d} | Train Loss: {train_loss:.4f} | "
          f"Train AUC: {train_auc:.4f} | Val AUC: {val_auc:.4f} | Best: {best_val_auc:.4f}")

    if wait >= PATIENCE:
        print(f"Early stopping at epoch {epoch}")
        break

resnet.load_state_dict(best_state)

🤔 Reflection 6.1 — Transfer Learning#

  1. We froze all layers except layer4 and fc. Why freeze early layers? What would happen if we fine-tuned all layers with a large learning rate?

  2. ResNet-18 was trained on natural images (dogs, cars, buildings). Why do its features transfer to chest X-rays? What is fundamentally similar about these image types?

  3. We upsampled our 28x28 images to 224x224 for ResNet. This doesn’t add information (you can’t create detail that wasn’t in the original). So why does ResNet still benefit from pre-training even when the input is low-resolution?

  4. We used only ~2.6M trainable parameters (layer4 + fc) out of 11.2M total. This is analogous to which approach in Lab 7? (Hint: think about DistilBERT fine-tuning.)

# == SOLUTION -- Reflection 6.1 ===================================================
# 1. Early layers learn universal low-level features (edges, textures, colors) that
#    are useful across all image types. These don't need to change for medical images.
#    Fine-tuning them with a large learning rate would cause "catastrophic forgetting"
#    -- destroying these valuable pre-trained features. Freezing them preserves the
#    general visual features and only adapts the higher-level, task-specific layers.
#    This is exactly the same principle as using a small learning rate for fine-tuning
#    DistilBERT (Lab 7).
#
# 2. All images (natural and medical) share low-level visual structure: edges, textures,
#    gradients, and spatial patterns. A horizontal edge in a lung boundary is detected
#    by the same filter as a horizontal edge on a building. The HIGHER-LEVEL features
#    (layer3, layer4) differ more -- ImageNet detects "dog ears" while we need "lung
#    consolidation" -- which is why we fine-tune the later layers.
#
# 3. Pre-training provides good INITIAL WEIGHTS and learned COMPUTATIONAL patterns
#    (how to compose features hierarchically). Even with low-resolution input, the
#    architecture's learned feature hierarchy (edges -> textures -> parts -> objects)
#    is valuable. The ResNet also benefits from its batch normalization statistics
#    and the residual connection structure, both learned during pre-training.
#    Fine-tuning adapts these to the specific input distribution.
#
# 4. This is directly analogous to DistilBERT fine-tuning: freeze most of the pre-
#    trained model and only update the task-specific layers. In DistilBERT, we
#    fine-tuned all layers but with a tiny learning rate (2e-5); here we freeze
#    early layers entirely and use a moderate learning rate (1e-4) for the rest.
#    Both strategies prevent catastrophic forgetting while adapting to the new task.
#    The principle is the same: leverage pre-trained representations, adjust minimally.
print("See comments above for solution.")

Part 7 — Model Comparison and Final Test Evaluation#

We now compare all approaches on the held-out test set. As in previous labs, this is the first time we touch the test set.

# -- Collect all validation results -------------------------------------------------
print("=== Validation Set Results (for model selection) ===\n")

all_val_results = {}

for name, res in flat_results.items():
    all_val_results[name] = res['val_auc']
for name, res in hog_results.items():
    all_val_results[name] = res['val_auc']

cnn_val_auc, _, _, _ = evaluate_cnn(cnn_model, val_dl)
all_val_results['CNN (from scratch)'] = cnn_val_auc

resnet_val_auc, _, _, _ = evaluate_resnet(resnet, val_dl_resnet)
all_val_results['ResNet-18 (transfer)'] = resnet_val_auc

for name, auc in all_val_results.items():
    print(f"  {name:30s}  Val AUROC = {auc:.4f}")

# Bar chart
fig, ax = plt.subplots(figsize=(12, 5))
names = list(all_val_results.keys())
aucs = list(all_val_results.values())
n_flat = len(flat_results)
n_hog = len(hog_results)
colors = (['#bdc3c7'] * n_flat + ['#95a5a6'] * n_hog + ['#3498db', '#e74c3c'])
bars = ax.barh(names, aucs, color=colors, edgecolor='white')
ax.set_xlabel('Validation AUROC')
ax.set_title('Model Comparison -- Validation AUROC')
ax.set_xlim(0.5, 1.0)
for bar, v in zip(bars, aucs):
    ax.text(v + 0.005, bar.get_y() + bar.get_height()/2, f'{v:.4f}',
            va='center', fontsize=10)
plt.tight_layout()
plt.show()

# -- Final test evaluation ---------------------------------------------------------
print("=== FINAL TEST SET RESULTS ===")
print("(All models selected on validation set; test set never seen before)\n")

# Best HOG model
best_hog_name = max(hog_results, key=lambda k: hog_results[k]['val_auc'])
best_hog_model = hog_results[best_hog_name]['model']
hog_test_proba = best_hog_model.predict_proba(
    scaler_hog.transform(X_test_hog))[:, 1]
hog_test_auc = roc_auc_score(y_test, hog_test_proba)

# Best flat model
best_flat_name = max(flat_results, key=lambda k: flat_results[k]['val_auc'])
best_flat_model = flat_results[best_flat_name]['model']
flat_test_proba = best_flat_model.predict_proba(
    scaler_flat.transform(X_test_flat))[:, 1]
flat_test_auc = roc_auc_score(y_test, flat_test_proba)

# CNN
cnn_test_auc, _, cnn_test_preds, cnn_test_labels = evaluate_cnn(cnn_model, test_dl)

# ResNet
resnet_test_auc, _, resnet_test_preds, resnet_test_labels = evaluate_resnet(resnet, test_dl_resnet)

test_summary = {
    f'Best Pixels ({best_flat_name.split("+")[0].strip()})': flat_test_auc,
    f'Best HOG ({best_hog_name.split("+")[0].strip()})': hog_test_auc,
    'CNN (from scratch)': cnn_test_auc,
    'ResNet-18 (transfer)': resnet_test_auc,
}

for name, auc_val in test_summary.items():
    print(f"  {name:40s}  Test AUROC = {auc_val:.4f}")

# ROC curves
fig, ax = plt.subplots(figsize=(8, 6))

fpr, tpr, _ = roc_curve(y_test, flat_test_proba)
ax.plot(fpr, tpr, label=f'Flat Pixels (AUC={flat_test_auc:.3f})', color='#bdc3c7', lw=2)

fpr, tpr, _ = roc_curve(y_test, hog_test_proba)
ax.plot(fpr, tpr, label=f'HOG (AUC={hog_test_auc:.3f})', color='#95a5a6', lw=2)

fpr, tpr, _ = roc_curve(cnn_test_labels, cnn_test_preds)
ax.plot(fpr, tpr, label=f'CNN (AUC={cnn_test_auc:.3f})', color='#3498db', lw=2)

fpr, tpr, _ = roc_curve(resnet_test_labels, resnet_test_preds)
ax.plot(fpr, tpr, label=f'ResNet-18 (AUC={resnet_test_auc:.3f})', color='#e74c3c', lw=2)

ax.plot([0, 1], [0, 1], 'k--', lw=1, alpha=0.5)
ax.set_xlabel('False Positive Rate')
ax.set_ylabel('True Positive Rate')
ax.set_title('ROC Curves -- Test Set')
ax.legend(loc='lower right')
plt.tight_layout()
plt.show()

🤔 Reflection 7.1 — Comparing Approaches#

  1. How do the four approaches rank? Does the pattern (hand-crafted features < CNN < transfer learning) hold? Where are the biggest performance jumps?

  2. Compare the progression across Labs 6-8:

    • Graphs: Fingerprints \(\approx\) GNN (hand-crafted features very competitive)

    • Text: TF-IDF < LSTM < DistilBERT (pre-training dominates)

    • Images: HOG < CNN < ResNet (where does the biggest gap lie?) What explains the different relative importance of hand-crafted vs. learned features across these modalities?

  3. Our dataset has only ~5K training images at 28x28 resolution. In real clinical imaging (e.g., CheXpert with 224K images at 320x320), the gap between CNN and transfer learning is often smaller. Why?

  4. If you were deploying a pneumonia detection system in a rural clinic with limited internet, no GPU, and need for interpretability, which approach would you choose and why?

# == SOLUTION -- Reflection 7.1 ===================================================
# 1. The typical ranking is: flat pixels < HOG < CNN < ResNet. The biggest jump
#    is usually from HOG to CNN (moving from hand-crafted to learned spatial
#    features). Transfer learning provides an additional boost, especially on small
#    datasets where training from scratch is data-limited. The gap between flat
#    pixels and HOG shows that SOME spatial feature engineering helps; the gap from
#    HOG to CNN shows that LEARNED spatial features are better.
#
# 2. The key factor is how well-designed the hand-crafted features are relative to
#    the task complexity. For MOLECULES: chemists designed fingerprints over decades,
#    and molecular structure is relatively simple -- so fingerprints are very good.
#    For TEXT: TF-IDF captures word importance well, but language understanding
#    requires compositionality and context that TF-IDF can't capture -- pre-training
#    provides this. For IMAGES: HOG captures local edges but misses the hierarchical
#    composition (edges -> textures -> parts -> objects) that CNNs learn. The more
#    complex the required feature hierarchy, the more learned features help.
#
# 3. With more data, CNNs trained from scratch have enough examples to learn good
#    feature hierarchies without pre-training. Transfer learning helps most when
#    data is scarce (it provides a good initialization). With 224K images, the CNN
#    can learn robust features on its own. The gap closes further at higher
#    resolutions where there's more signal per image.
#
# 4. For a rural clinic: HOG + Logistic Regression. Reasons: (a) runs on CPU in
#    milliseconds (no GPU needed), (b) model is tiny (~1KB) and doesn't need
#    internet after deployment, (c) coefficients are somewhat interpretable
#    (which HOG cells contribute most), (d) easy to validate and audit. If
#    slightly more resources are available, a small CNN (~170K params) could run
#    on a smartphone. ResNet is likely overkill for this setting and adds
#    deployment complexity without proportional clinical benefit.
print("See comments above for solution.")

Part 8 — Extensions: What Can You Do From Here?#

If you have extra time or want to explore further, try any of these extensions. They are optional and not graded.

# -- Extension 1: Data Augmentation -------------------------------------------------
# Train the CNN with and without data augmentation to see the effect

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
])

# Augmented training loop
print("Training CNN WITH data augmentation...")
cnn_aug = SimpleCNN(dropout=0.3).to(device)
optimizer_aug = torch.optim.Adam(cnn_aug.parameters(), lr=1e-3)
criterion = nn.BCELoss()

aug_history = {'val_auc': []}
for epoch in range(1, 21):
    cnn_aug.train()
    for X, y_batch in train_dl:
        X, y_batch = X.to(device), y_batch.to(device)
        # Apply augmentation
        X_aug = train_transform(X)
        optimizer_aug.zero_grad()
        out = cnn_aug(X_aug)
        loss = criterion(out, y_batch)
        loss.backward()
        optimizer_aug.step()

    val_auc, _, _, _ = evaluate_cnn(cnn_aug, val_dl)
    aug_history['val_auc'].append(val_auc)
    if epoch % 5 == 0:
        print(f"  Epoch {epoch:2d} | Val AUC: {val_auc:.4f}")

# Compare
plt.figure(figsize=(8, 4))
plt.plot(cnn_history['val_auc'], label='No augmentation', color='#e74c3c', lw=2)
plt.plot(aug_history['val_auc'], label='With augmentation', color='#3498db', lw=2)
plt.xlabel('Epoch'); plt.ylabel('Validation AUROC')
plt.title('Effect of Data Augmentation')
plt.legend(); plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

🧠 Final Reflection — Image Data in the ML Landscape#

Now that you’ve worked through images alongside tabular (Labs 0-5), graph (Lab 12), and text (Lab 13) data, answer these synthesis questions:

  1. The modality progression: Across Labs 12-14, we consistently saw a pattern: tabularization baseline \(\to\) domain-specific learned model \(\to\) pre-trained + fine-tuned model. In which modality was the gap between stages largest? What does this tell you about the nature of each data type?

  2. Inductive biases summary: We’ve now seen four architectures for structured data: GNNs (message passing over graphs), LSTMs (sequential processing of text), CNNs (local filters over grids), and Transformers (global self-attention). For each architecture, state the key assumption about data structure and give a medical data type where that assumption holds.

  3. Transfer learning across modalities: Transfer learning from ImageNet helped for images, and from Wikipedia/BookCorpus helped for text. Why is transfer learning harder for graph data? What would a “foundation model” for molecular graphs look like?

  4. The deployment reality: In clinical practice, model accuracy is necessary but not sufficient. A deployed pneumonia detection system must also handle: (a) distribution shift (X-rays from different machines), (b) calibration (predicted probabilities must be meaningful), (c) failure modes (model must know when it’s uncertain), (d) bias (fair performance across demographics). Which of these is most challenging for CNNs specifically?

  5. Cross-modal connections: Modern health AI increasingly uses multi-modal models that combine images, text, and structured data (e.g., radiology images + clinical notes + lab values). How would you design a model that takes BOTH a chest X-ray and the clinical note as input? What are the key architectural decisions?

# == SOLUTION -- Final Reflection ===================================================
# 1. The gap between stages varies by modality:
#    - GRAPHS: small gap between fingerprints and GNNs (well-designed features capture
#      most signal in small molecular datasets)
#    - TEXT: large gap between TF-IDF and DistilBERT (language understanding requires
#      compositionality that only pre-training provides at scale)
#    - IMAGES: moderate gap -- HOG captures local patterns well, but CNNs capture
#      hierarchical composition that HOG misses; transfer learning adds more for small
#      datasets. The lesson: the more complex the required feature hierarchy, the more
#      learned + pre-trained features help.
#
# 2. Architecture -> Assumption -> Medical data:
#    - GNN: neighbors influence each other (graph topology matters) -> molecular graphs,
#      protein structures, patient similarity networks
#    - LSTM: sequential order matters -> clinical notes, temporal lab values, medication
#      sequences
#    - CNN: local spatial patterns matter (translation invariance) -> X-rays, pathology
#      slides, dermoscopy images
#    - Transformer: any token can attend to any other -> clinical documents (long-range
#      context), multi-modal fusion, genomic sequences
#
# 3. Transfer learning is harder for graphs because: (a) there's no single "ImageNet
#    for molecules" -- molecular properties are diverse and task-specific, (b) graph
#    structures vary enormously (small drug molecules vs. large proteins vs. social
#    networks), (c) node features differ across domains. A molecular foundation model
#    would pre-train on millions of unlabeled molecular structures using self-supervised
#    objectives (predict masked atoms, predict 3D coordinates, contrastive learning on
#    molecular substructures). Models like Uni-Mol and MolBERT are early examples.
#
# 4. Distribution shift is most challenging for CNNs because they learn subtle pixel-
#    level patterns that are highly scanner-specific. A model trained on X-rays from
#    GE scanners may fail on Siemens scanners because of different image characteristics
#    (contrast, noise, resolution). This is less of a problem for text models (clinical
#    language is more standardized than imaging hardware). Mitigation: train on diverse
#    sources, use domain adaptation techniques, or add scanner/site as an input feature.
#
# 5. Multi-modal architecture: (a) encode each modality with its appropriate backbone
#    (CNN for X-ray, Transformer for clinical note, MLP for lab values), (b) project
#    each into a shared embedding space of the same dimension, (c) fuse using cross-
#    attention (image features attend to text, and vice versa) or simple concatenation,
#    (d) classify from the fused representation. Key decisions: (i) early vs late
#    fusion (combine raw features vs combine final embeddings), (ii) cross-attention
#    vs concatenation (richer interaction vs simpler training), (iii) how to handle
#    missing modalities (not all patients have both X-ray and notes).
print("See comments above for solution.")