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')

GPU is available, using GPU Tesla T4

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

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))

Training dataset has length 48000
Training dataset has length 12000
Training dataset has length 10000

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

Input has shape torch.Size([1, 28, 28]) and the output is 7

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(
    nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    # nn.Dropout(p=0.5), # Zero out 50% of all pixels
    # nn.Dropout2d(p=0.5), # Zero out 50% of all channels
)

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

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

Model transforms input from size torch.Size([1, 1, 28, 28]) into torch.Size([1, 32, 7, 7])

Sequential(
  (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (1): ReLU()
  (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (4): ReLU()
  (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)

model.add_module('flatten', nn.Flatten())

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

Model transforms input from size torch.Size([1, 1, 28, 28]) into torch.Size([1, 1568])

Sequential(
  (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (1): ReLU()
  (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (4): ReLU()
  (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (flatten): Flatten(start_dim=1, end_dim=-1)
)

model.add_module('dropout', nn.Dropout(p=0.5)) # We can also just apply dropout here to improve generalization
model.add_module('fc1', nn.Linear(1568, 256))
model.add_module('relu', nn.ReLU())
model.add_module('fc2', nn.Linear(256, 10))

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

Model transforms input from size torch.Size([1, 1, 28, 28]) into torch.Size([1, 10])

Sequential(
  (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (1): ReLU()
  (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (4): ReLU()
  (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (dropout): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=1568, out_features=256, bias=True)
  (relu): ReLU()
  (fc2): Linear(in_features=256, out_features=10, bias=True)
)

def make_model():
    model = nn.Sequential(
        nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding='same'),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding='same'),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2)
    )
    model.add_module('flatten', nn.Flatten())
    model.add_module('dropout', nn.Dropout(p=0.5))
    model.add_module('fc1', nn.Linear(1568, 256))
    model.add_module('relu', nn.ReLU())
    model.add_module('fc2', nn.Linear(256, 10))
    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 = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)

    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.
            x = x.to(device)
            y = y.to(device)

            # Make a prediction, compute the loss
            y_pred = model(x)
            loss = loss_fn(y_pred, y)

            # Backpropagate gradients, apply weight updates, reset gradients
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            # 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.
                x = x.to(device)
                y = y.to(device)
    
                # Make a prediction, compute the loss
                y_pred = model(x)
                loss = loss_fn(y_pred, y)
        
                # 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)

model = make_model()
loss = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model, model_stats = train(
    model, 
    loss, 
    optimizer, 
    mnist_train, 
    mnist_val, 
    num_epochs=10,
    batch_size=64,
    num_workers=4,
)

Epoch 1 | Train loss 0.258 | Train acc 0.918 | Val loss 0.083 | Val acc 0.975

Epoch 2 | Train loss 0.089 | Train acc 0.972 | Val loss 0.051 | Val acc 0.984

Epoch 3 | Train loss 0.066 | Train acc 0.978 | Val loss 0.042 | Val acc 0.987

Epoch 4 | Train loss 0.055 | Train acc 0.983 | Val loss 0.034 | Val acc 0.990

Epoch 5 | Train loss 0.048 | Train acc 0.985 | Val loss 0.033 | Val acc 0.990

Epoch 6 | Train loss 0.043 | Train acc 0.986 | Val loss 0.032 | Val acc 0.989

Epoch 7 | Train loss 0.037 | Train acc 0.988 | Val loss 0.030 | Val acc 0.990

Epoch 8 | Train loss 0.033 | Train acc 0.990 | Val loss 0.031 | Val acc 0.991

Epoch 9 | Train loss 0.031 | Train acc 0.990 | Val loss 0.033 | Val acc 0.990

Epoch 10 | Train loss 0.029 | Train acc 0.990 | Val loss 0.029 | Val acc 0.992

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):
        xb = xb.to(device)
        yb = yb.to(device)
        y_pred = model(xb)
        y_pred = torch.argmax(y_pred, dim=1)
        test_correct_ += (y_pred == yb).float().sum().item()
    test_acc = test_correct_ / len(test_loader.dataset)

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

Test accuracy: 0.992

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.detach().cpu().numpy()
print(f'Feature kernels have shape {feature_kernels.shape} and '
      f'values between {feature_kernels.min():.3f} and {feature_kernels.max():.3f}')

Feature kernels have shape (32, 1, 3, 3) and values between -0.655 and 0.627

fig, axes = plt.subplots(4, 8, figsize=(10, 5), dpi=150)

for kernel, ax in zip(feature_kernels, axes.ravel()):
    ax.imshow(kernel.squeeze(), cmap='bwr', vmin=-1, vmax=1)
    ax.set(xticks=[], yticks=[])

CNN classification with a more complicated dataset¶

Dataset: STL10¶

import torch
import torchvision

# Set download=True if you haven't accessed this dataset before
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')

# Set download=True if you haven't accessed this dataset before
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')

Train/val dataset has size (5000, 3, 96, 96) and targets (5000,)
Train/val dataset has 10 unique labels
Test dataset has size (8000, 3, 96, 96) and targets (8000,)
Test dataset has 10 unique labels

Another dataset: CIFAR10¶

import torch
import torchvision

# Set download=True if you haven't accessed this dataset before
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
# Expect better accuracy than what we got for STL10
img_train_val_dataset.labels = np.array(img_train_val_dataset.targets)

# Set download=True if you haven't accessed this dataset before
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')

Train/val dataset has size (50000, 32, 32, 3) and targets 50000
Train/val dataset has 10 unique labels
Test dataset has size (10000, 32, 32, 3) and targets 10000
Test dataset has 10 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)

Image type: <class 'torchvision.tv_tensors._image.Image'>
Image shape: torch.Size([3, 96, 96])
Image min: 0.0
Image max: 1.0
Label: 1

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
img_hflip = F.hflip(img)
ax[1].imshow(img_hflip.permute(1, 2, 0))
ax[1].set_title('Horizontal flip')

# Vertical flip
img_vflip = F.vflip(img)
ax[2].imshow(img_vflip.permute(1, 2, 0))
ax[2].set_title('Vertical flip')

# Image rotation
img_rot = F.rotate(img, 45)
ax[3].imshow(img_rot.permute(1, 2, 0))
ax[3].set_title('45$^{\\circ}$ rotation')

# Adjust brightness
img_bri = F.adjust_brightness(img, brightness_factor=0.6)
ax[4].imshow(img_bri.permute(1, 2, 0))
ax[4].set_title('Adjust brightness')

# Adjust contrast
img_con = F.adjust_contrast(img, contrast_factor=5)
ax[5].imshow(img_con.permute(1, 2, 0))
ax[5].set_title('Adjust contrast')

# Blur
img_blr = F.gaussian_blur(img, kernel_size=5)
ax[6].imshow(img_blr.permute(1, 2, 0))
ax[6].set_title('Blurring')

# Noise
img_nse = F.permute_channels(img, (2, 0, 1))
ax[7].imshow(img_nse.permute(1, 2, 0))
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(),
    v2.RandomHorizontalFlip(p=0.5),
    # v2.RandomVerticalFlip(p=0.5),
    v2.RandomRotation(degrees=[-30, 30]),
    v2.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
    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()

Images has shape torch.Size([32, 3, 96, 96])
Labels has shape torch.Size([32])

CNN definition¶

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

cnn_base = nn.Sequential(
    # [B, 3, 96, 96]
    nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Dropout(p=0.2),
    # [B, 32, 48, 48]
    nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Dropout(p=0.2),
    # [B, 64, 24, 24]
    nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Dropout(p=0.2),
    # [B, 128, 12, 12]
)

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

CNN has 93248 trainable parameters
Model output has size: torch.Size([32, 128, 12, 12])

Sequential(
  (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (1): ReLU()
  (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (3): Dropout(p=0.2, inplace=False)
  (4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (5): ReLU()
  (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (7): Dropout(p=0.2, inplace=False)
  (8): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
  (9): ReLU()
  (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (11): Dropout(p=0.2, inplace=False)
)

Different “classifier heads”¶

global_max_pool_classifier = nn.Sequential(
    # Input [B, 128, 12, 12]
    nn.MaxPool2d(kernel_size=12),# STL10
    # nn.MaxPool2d(kernel_size=3), #CIFAR10
    # [B, 128, 1, 1]
    nn.Flatten(),
    # [B, 128]
    nn.Linear(128, 10)
)

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

CNN has 94538 trainable parameters
Model output has size: torch.Size([32, 10])

Sequential(
  (0): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Dropout(p=0.2, inplace=False)
    (4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (5): ReLU()
    (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (7): Dropout(p=0.2, inplace=False)
    (8): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (9): ReLU()
    (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (11): Dropout(p=0.2, inplace=False)
  )
  (1): Sequential(
    (0): MaxPool2d(kernel_size=12, stride=12, padding=0, dilation=1, ceil_mode=False)
    (1): Flatten(start_dim=1, end_dim=-1)
    (2): Linear(in_features=128, out_features=10, bias=True)
  )
)

global_avg_pool_classifier = nn.Sequential(
    # Input [B, 128, 12, 12]
    nn.AvgPool2d(kernel_size=12),# STL10
    # nn.AvgPool2d(kernel_size=3), #CIFAR10
    # [B, 128, 1, 1]
    nn.Flatten(),
    # [B, 128]
    nn.Linear(128, 10)
)

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

CNN has 94538 trainable parameters
Model output has size: torch.Size([32, 10])

Sequential(
  (0): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Dropout(p=0.2, inplace=False)
    (4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (5): ReLU()
    (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (7): Dropout(p=0.2, inplace=False)
    (8): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (9): ReLU()
    (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (11): Dropout(p=0.2, inplace=False)
  )
  (1): Sequential(
    (0): AvgPool2d(kernel_size=12, stride=12, padding=0)
    (1): Flatten(start_dim=1, end_dim=-1)
    (2): Linear(in_features=128, out_features=10, bias=True)
  )
)

fully_connected_classifier = nn.Sequential(
    # Input [B, 128, 12, 12]
    nn.Flatten(),
    # [B, 128 * 12 * 12]
    nn.Linear(128 * 12 * 12, 128),
    nn.ReLU(),
    # [B, 128]
    nn.Linear(128, 10)
)

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

CNN has 2453962 trainable parameters
Model output has size: torch.Size([32, 10])

Sequential(
  (0): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Dropout(p=0.2, inplace=False)
    (4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (5): ReLU()
    (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (7): Dropout(p=0.2, inplace=False)
    (8): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=same)
    (9): ReLU()
    (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (11): Dropout(p=0.2, inplace=False)
  )
  (1): Sequential(
    (0): Flatten(start_dim=1, end_dim=-1)
    (1): Linear(in_features=18432, out_features=128, bias=True)
    (2): ReLU()
    (3): Linear(in_features=128, out_features=10, bias=True)
  )
)

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')

MLP has 3540362 trainable parameters

CNN training¶

# Define a loss function and optimizer
loss = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model, model_stats = train(
    model, 
    loss, 
    optimizer, 
    img_train_dataset, 
    img_val_dataset, 
    num_epochs=25,
    batch_size=64,
    num_workers=4,
)

Epoch 1 | Train loss 2.089 | Train acc 0.222 | Val loss 1.818 | Val acc 0.354

Epoch 2 | Train loss 1.725 | Train acc 0.348 | Val loss 1.647 | Val acc 0.368

Epoch 3 | Train loss 1.570 | Train acc 0.417 | Val loss 1.569 | Val acc 0.423

Epoch 4 | Train loss 1.527 | Train acc 0.429 | Val loss 1.462 | Val acc 0.462

Epoch 5 | Train loss 1.426 | Train acc 0.480 | Val loss 1.442 | Val acc 0.451

Epoch 6 | Train loss 1.343 | Train acc 0.507 | Val loss 1.366 | Val acc 0.508

Epoch 7 | Train loss 1.322 | Train acc 0.508 | Val loss 1.293 | Val acc 0.531

Epoch 8 | Train loss 1.231 | Train acc 0.549 | Val loss 1.259 | Val acc 0.524

Epoch 9 | Train loss 1.198 | Train acc 0.558 | Val loss 1.243 | Val acc 0.539

Epoch 10 | Train loss 1.163 | Train acc 0.571 | Val loss 1.269 | Val acc 0.536

Epoch 11 | Train loss 1.113 | Train acc 0.592 | Val loss 1.172 | Val acc 0.589

Epoch 12 | Train loss 1.061 | Train acc 0.613 | Val loss 1.183 | Val acc 0.575

Epoch 13 | Train loss 1.042 | Train acc 0.617 | Val loss 1.207 | Val acc 0.557

Epoch 14 | Train loss 0.992 | Train acc 0.635 | Val loss 1.204 | Val acc 0.575

Epoch 15 | Train loss 1.000 | Train acc 0.633 | Val loss 1.169 | Val acc 0.572

Epoch 16 | Train loss 0.997 | Train acc 0.644 | Val loss 1.162 | Val acc 0.587

Epoch 17 | Train loss 0.941 | Train acc 0.659 | Val loss 1.197 | Val acc 0.572

Epoch 18 | Train loss 0.892 | Train acc 0.673 | Val loss 1.279 | Val acc 0.567

Epoch 19 | Train loss 0.894 | Train acc 0.676 | Val loss 1.148 | Val acc 0.589

Epoch 20 | Train loss 0.848 | Train acc 0.700 | Val loss 1.148 | Val acc 0.597

Epoch 21 | Train loss 0.828 | Train acc 0.701 | Val loss 1.136 | Val acc 0.616

Epoch 22 | Train loss 0.808 | Train acc 0.703 | Val loss 1.218 | Val acc 0.582

Epoch 23 | Train loss 0.795 | Train acc 0.712 | Val loss 1.112 | Val acc 0.585

Epoch 24 | Train loss 0.761 | Train acc 0.724 | Val loss 1.188 | Val acc 0.595

Epoch 25 | Train loss 0.746 | Train acc 0.724 | Val loss 1.164 | Val acc 0.590

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):
        xb = xb.to(device)
        yb = yb.to(device)
        y_pred = model(xb)
        y_pred = torch.argmax(y_pred, dim=1)
        test_correct_ += (y_pred == yb).float().sum().item()
    test_acc = test_correct_ / len(test_loader.dataset)

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

Test accuracy: 0.609

# 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()

torch.Size([64, 3, 96, 96]) torch.Size([64])

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¶