CNN classification

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
plt.rcParams['xtick.direction'] = 'in'
plt.rcParams['ytick.direction'] = 'in'
plt.rcParams['axes.grid'] = True
plt.rcParams['axes.axisbelow'] = True
plt.rcParams['figure.dpi'] = 150

# Check available hardware
import torch
if torch.cuda.is_available():
    print(f'GPU is available, using GPU {torch.cuda.get_device_name(0)}')
    device = torch.device('cuda')
elif torch.mps.is_available():
    print(f'Apple Metal Performance Shaders framework is available, using {torch.device("mps")}')
    device = torch.device('mps')
else:
    print(f'Hardware acceleration not available, using CPU')
    device = torch.device('cpu')

Building a CNN¶

Now that we understand convolution, and why we should care about it, we are almost ready to implement a CNN. In general, a CNN architecture contains:

Convolution layers which execute a discrete convolution with a learnable weighting or kernel function
Pooling layers which perform spatial downsampling (mimicking the role of complex cells in a visual cortex)
Fully-connected layers which make determinations based on the high-level features at the end of the network

Dataset: MNIST¶

For this example, we’ll use MNIST. But since we’re using torch now, we can use the MNIST located in torchvision

import torch
import torchvision
from torchvision.transforms import v2

# We'll talk more about this later in the demo
transform = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True)
])
mnist = torchvision.datasets.MNIST(root='~/scr10/', train=True, transform=transform, download=True)
mnist_test = torchvision.datasets.MNIST(root='~/scr10/', train=False, transform=transform, download=False)
from torch.utils.data import random_split
mnist_train, mnist_val = random_split(mnist, [0.8, 0.2])
print('Training dataset has length', len(mnist_train))
print('Training dataset has length', len(mnist_val))
print('Training dataset has length', len(mnist_test))

x, y = mnist_train[0]
print('Input has shape', x.shape, 'and the output is', y)

fig, ax = plt.subplots(2, 5, sharex=True, sharey=True)
ax = ax.flatten()
for ii in range(10):
    x, y = mnist_train[np.random.choice(len(mnist_train))]
    ax[ii].imshow(x.squeeze(), cmap='Greys')
    ax[ii].set(xticks=[], yticks=[])
plt.tight_layout()

CNN definition¶

import torch.nn as nn
model = nn.Sequential(
    # Your CNN layers go here
)

z = model(x[None])
print(f'Model transforms input from size {x[None].shape} into {z.shape}')

# Visualizing some feature maps
fig, axes = plt.subplots(4, 8, figsize=(10, 5), dpi=150)
for feature, ax in zip(z[0], axes.ravel()):
    ax.imshow(feature.detach().numpy(), cmap='Greys', vmin=0, vmax=1)
    ax.set(xticks=[], yticks=[])

model

# Flatten() converts spatial feature maps to a more typical hidden layer format

z = model(x[None])
print(f'Model transforms input from size {x[None].shape} into {z.shape}')
model

# Classifier head translates model outputs into predictions

z = model(x[None])
print(f'Model transforms input from size {x[None].shape} into {z.shape}')
model

def make_model():
    """ Function to make a new model, so we don't accidentally retrain
        a model that we've already trained
    """
    model = # Your code here
    return model

CNN training¶

from tqdm.notebook import tqdm, trange
from torch.utils.data import DataLoader

def train(model, loss_fn, optimizer, train_dataset, val_dataset, 
          batch_size=32, num_epochs=100, num_workers=0, device=device):
    # Send model to correct hardware device
    model = model.to(device)

    # Create dataloaders for train/validation
    train_loader = 
    val_loader = 

    train_loss, train_acc = [], []
    val_loss, val_acc = [], []
    
    for epoch in range(num_epochs):
        train_loss_ = 0.
        train_acc_ = 0.
        model.train()
        for x, y, in tqdm(train_loader, leave=False):
            # Send x, y to GPU/TPU/etc.

            # Make a prediction, compute the loss

            # Backpropagate gradients, apply weight updates, reset gradients

            # Logging
            train_loss_ += loss.item() * x.shape[0]
            train_acc_  += (torch.argmax(y_pred, dim=1) == y).float().sum().item()

        train_loss.append(train_loss_ / len(train_dataset))
        train_acc.append(train_acc_ / len(train_dataset))

        val_loss_ = 0.
        val_acc_ = 0.
        model.eval()
        with torch.no_grad():
            for x, y, in tqdm(val_loader, leave=False):
                # Send x, y to GPU/TPU/etc.
    
                # Make a prediction, compute the loss
        
                # Logging
                val_loss_ += loss.item() * x.shape[0]
                val_acc_  += (torch.argmax(y_pred, dim=1) == y).float().sum().item()

            val_loss.append(val_loss_ / len(val_dataset))
            val_acc.append(val_acc_ / len(val_dataset))

        print(f'Epoch {epoch+1} | '
              f'Train loss { train_loss[-1]:.3f} | '
              f'Train acc { train_acc[-1]:.3f} | '
              f'Val loss { val_loss[-1]:.3f} | '
              f'Val acc { val_acc[-1]:.3f}'
        )

    return model, (train_loss, train_acc, val_loss, val_acc)

# Initialize the model
model = make_model()

# Initialize a loss function and optimizer

model, model_stats = train(
    model, 
    loss, 
    optimizer, 
    mnist_train, 
    mnist_val, 
    num_epochs=10,
    batch_size=64,
    num_workers=4,
)

# Plot loss curves
fig, ax = plt.subplots(1, 2, sharex=True, figsize=(5, 2.5))

train_loss, train_acc, val_loss, val_acc = model_stats
ax[0].plot(train_loss, color='black', label='Train')
ax[0].plot(val_loss, color='red', label='Validation')
ax[0].set(xlabel='Epoch', ylabel='Loss', yscale='log')

ax[1].plot(train_acc, color='black', label='Train')
ax[1].plot(val_acc, color='red', label='Validation')
ax[1].set(xlabel='Epoch', ylabel='Accuracy')

fig.legend(*ax[1].get_legend_handles_labels(), 
           loc='center left', bbox_to_anchor=[1, 0.5], framealpha=0)

plt.tight_layout()

Model inference¶

# Inference
model.eval()
test_loader = DataLoader(mnist_test, batch_size=128, shuffle=False)
with torch.no_grad():
    test_correct_ = 0.
    for xb, yb in tqdm(test_loader,leave=False):
        # Send x, y to device
        xb = xb.to(device)
        yb = yb.to(device)

        # Make a prediction
        y_pred = model(xb)
        y_pred = torch.argmax(y_pred, dim=1)

        # Log the result
        test_correct_ += (y_pred == yb).float().sum().item()
    test_acc = test_correct_ / len(test_loader.dataset)

print(f"Test accuracy: {test_acc:.3f}")

fig, ax = plt.subplots(2, 5, figsize=(10, 4), sharex=True, sharey=True)
ax = ax.flatten()
for ii in range(10):
    ax[ii].imshow(xb[ii].squeeze().cpu().numpy(), cmap='Greys')
    ax[ii].text(0.5, -0.2, f'True: {yb[ii]}', va='bottom', ha='center', transform=ax[ii].transAxes)
    ax[ii].text(0.5, -0.2, f'Pred: {y_pred[ii]}', color='red', va='top', ha='center', transform=ax[ii].transAxes)
    ax[ii].set(xticks=[], yticks=[])
plt.tight_layout()

Visualizing CNN activations¶

Check out this fun online tool for NN weight visualization

feature_kernels = model[0].weight

CNN classification with a more complicated dataset¶

Dataset: STL10¶

import torch
import torchvision

img_train_val_dataset = torchvision.datasets.STL10(split='train', root='~/scr10/', download=False)
print(f'Train/val dataset has size {img_train_val_dataset.data.shape} and targets {img_train_val_dataset.labels.shape}')
print(f'Train/val dataset has {len(img_train_val_dataset.classes)} unique labels')

img_test_dataset = torchvision.datasets.STL10(split='test', root='~/scr10/', download=False)
print(f'Test dataset has size {img_test_dataset.data.shape} and targets {img_test_dataset.labels.shape}')
print(f'Test dataset has {len(img_test_dataset.classes)} unique labels')

Another dataset: CIFAR10¶

import torch
import torchvision

img_train_val_dataset = torchvision.datasets.CIFAR10(train=True, root='~/scr10/', download=False)
print(f'Train/val dataset has size {img_train_val_dataset.data.shape} and targets {len(img_train_val_dataset.targets)}')
print(f'Train/val dataset has {len(img_train_val_dataset.classes)} unique labels')

# I developed this notebook for STL10, so add this line if you want compatibility for CIFAR10
img_train_val_dataset.labels = np.array(img_train_val_dataset.targets)

img_test_dataset = torchvision.datasets.CIFAR10(train=False, root='~/scr10/', download=False)
print(f'Test dataset has size {img_test_dataset.data.shape} and targets {len(img_test_dataset.targets)}')
print(f'Test dataset has {len(img_test_dataset.classes)} unique labels')

Visualization¶

fig, ax = plt.subplots(2, 5, sharex=True, sharey=True)
ax = ax.flatten()
for ii in range(10):
    class_idx = np.argwhere(img_train_val_dataset.labels == ii)
    idx = class_idx[0,0]
    img, label = img_train_val_dataset[idx]
    ax[ii].imshow(img)
    ax[ii].axis('off')
    ax[ii].set_title(img_train_val_dataset.classes[ii])

plt.tight_layout()

from sklearn.decomposition import PCA
pca = PCA(n_components=2, whiten=True)

X_in = img_train_val_dataset.data.reshape([len(img_train_val_dataset), -1])
y_in = img_train_val_dataset.labels
X_pca = pca.fit_transform(X_in)

ax = plt.figure().gca()
for ii in range(10):
    mask = y_in == ii
    ax.scatter(X_pca[mask, 0], X_pca[mask, 1], s=5, label=img_train_val_dataset.classes[ii])
ax.set(xlabel='Principal Component 1', ylabel='Principal Component 2')
ax.legend()
plt.tight_layout()

Image transforms¶

We don’t have that much training data, only 5,000 images. So we apply transformations to augment our dataset. This can do several things. Like SGD, it may introduce stochasticity to improve model training and avoid local minima. But it may also improve our model’s ability to generalize by making it more robust to particular transformations

from torchvision.transforms import v2

# The most basic transformation turns a PIL image into a torch tensor
transform = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True)
])
img_train_val_dataset.transform = transform
img, label = img_train_val_dataset[0]
print('Image type:', type(img))
print('Image shape:', img.shape)
print('Image min:', img.min().item())
print('Image max:', img.max().item()) 
print('Label:', label)

from torchvision.transforms.v2 import functional as F

fig, ax = plt.subplots(2, 4, figsize=(8, 4))
ax = ax.flatten()

img, label = img_train_val_dataset[0]
ax[0].imshow(img.permute(1, 2, 0))
ax[0].set_title('Basic image')

# Horizontal flip
ax[1].set_title('Horizontal flip')

# Vertical flip
ax[2].set_title('Vertical flip')

# Image rotation
ax[3].set_title('45$^{\\circ}$ rotation')

# Adjust brightness
ax[4].set_title('Adjust brightness')

# Adjust contrast
ax[5].set_title('Adjust contrast')

# Blur
ax[6].set_title('Blurring')

# Noise
ax[7].set_title('Color permutation')

for a in ax:
    a.axis('off')

Create train, validation, datasets, check augmentations¶

# Decide on our training transform
transform = v2.Compose([
    v2.ToImage(),
    # TorchVision Transforms go here
    # ...
    v2.ToDtype(torch.float32, scale=True)
])

from torch.utils.data import random_split, DataLoader

# Give the training dataset the random transformation
img_train_val_dataset.transform = transform
img_train_dataset, img_val_dataset = random_split(img_train_val_dataset, lengths=[0.8, 0.2])

batch_size=32
loader = DataLoader(img_train_dataset, batch_size=batch_size, shuffle=True)

for imgs, labels in loader:
    print('Images has shape', imgs.shape)
    print('Labels has shape', labels.shape)
    break

fig, axes = plt.subplots(4, 8, figsize=(10, 5))
for ax, img, label in zip(axes.flatten(), imgs, labels):
    ax.imshow(img.permute(1, 2, 0))
    ax.set_title(img_train_val_dataset.classes[label])
    ax.axis('off')

plt.tight_layout()

CNN definition¶

import torch.nn as nn
import torch.nn.functional as F

cnn_base = nn.Sequential(
    # Input: [B, 3, 96, 96]
    # Your CNN layers go here
)

num_trainable_params = sum(p.numel() for p in cnn_base.parameters() if p.requires_grad)
print(f'CNN has {num_trainable_params} trainable parameters')

# Check that the model can accept a batch of images
print(f'Model output has size:', cnn_base.forward(imgs).shape)

cnn_base

Different classifier heads¶

global_max_pool_classifier = nn.Sequential(
    # Your classifier layers here
)

model = nn.Sequential(
    cnn_base,
    global_max_pool_classifier
)

num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'CNN has {num_trainable_params} trainable parameters')

# Check that the model can accept a batch of images
print(f'Model output has size:', model.forward(imgs).shape)

model

global_avg_pool_classifier = nn.Sequential(
    # Your classifier layers here
)

model = nn.Sequential(
    cnn_base,
    global_avg_pool_classifier
)

num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'CNN has {num_trainable_params} trainable parameters')

# Check that the model can accept a batch of images
print(f'Model output has size:', model.forward(imgs).shape)

model

fully_connected_classifier = nn.Sequential(
    # Your classifier layers here
)

model = nn.Sequential(
    cnn_base,
    fully_connected_classifier
)

num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'CNN has {num_trainable_params} trainable parameters')

# Check that the model can accept a batch of images
print(f'Model output has size:', model.forward(imgs).shape)

model

The model above has fewer than 100,000 trainable parameters. This is pretty small! As a comparison, let’s look at a multi-layer perceptron which directly goes from the image to a 128-dimensional hidden layer, and from that to an output

mlp = nn.Sequential(
    nn.Flatten(),
    nn.Linear(3*96*96, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

num_trainable_params = sum(p.numel() for p in mlp.parameters() if p.requires_grad)
print(f'MLP has {num_trainable_params} trainable parameters')

CNN training¶

# Define a loss function and optimizer
loss = 
optimizer = 
model, model_stats = train(
    model, 
    loss, 
    optimizer, 
    img_train_dataset, 
    img_val_dataset, 
    num_epochs=25,
    batch_size=64,
    num_workers=4,
)

fig, ax = plt.subplots(1, 2, sharex=True, figsize=(5, 2.5))

train_loss, train_acc, val_loss, val_acc = model_stats
ax[0].plot(train_loss, color='black', label='Train')
ax[0].plot(val_loss, color='red', label='Validation')
ax[0].set(xlabel='Epoch', ylabel='Loss')

ax[1].plot(train_acc, color='black', label='Train')
ax[1].plot(val_acc, color='red', label='Validation')
ax[1].set(xlabel='Epoch', ylabel='Accuracy')

fig.legend(*ax[1].get_legend_handles_labels(), 
           loc='center left', bbox_to_anchor=[1, 0.5], framealpha=0)

plt.tight_layout()

Model inference¶

# Inference
transform = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True)
])
img_test_dataset.transform = transform

model.eval()
test_loader = DataLoader(img_test_dataset, batch_size=128, shuffle=False)
with torch.no_grad():
    test_correct_ = 0.
    for xb, yb in tqdm(test_loader,leave=False):
        # Send x, y to device
        xb = xb.to(device)
        yb = yb.to(device)

        # Make a prediction
        y_pred = model(xb)
        y_pred = torch.argmax(y_pred, dim=1)

        # Log the result
        test_correct_ += (y_pred == yb).float().sum().item()
    test_acc = test_correct_ / len(test_loader.dataset)

print(f"Test accuracy: {test_acc:.3f}")

# Look at 8 examples
fig, axes = plt.subplots(2, 4, figsize=(8, 4))
with torch.no_grad():
    xb = xb.to(device)
    y_pred = model(xb)
    y_pred = F.softmax(y_pred, dim=1)
    y_pred = y_pred.cpu().numpy()
for ax, img, label, pred in zip(axes.flatten(), xb, yb, y_pred):
    corr_prob = pred[label]
    pred_prob = pred.max()
    pred_label = np.argmax(pred)
    ax.imshow(img.permute(1, 2, 0).cpu())
    ax.set_title(
        f'GT: {img_test_dataset.classes[label]} ($p$ = {corr_prob:.2f})\n'
        f'Pred: {img_test_dataset.classes[pred_label]} ($p$ = {pred_prob:.2f})')
    ax.axis('off')

plt.tight_layout()

Building a CNN¶

Dataset: MNIST¶

CNN definition¶

CNN training¶

Model inference¶

Visualizing CNN activations¶

CNN classification with a more complicated dataset¶

Dataset: STL10¶

Another dataset: CIFAR10¶

Visualization¶

Image transforms¶

Create train, validation, datasets, check augmentations¶

CNN definition¶

Different classifier heads¶

Aside on weight sharing¶

CNN training¶

Model inference¶