# RUN THIS CELL FIRST.
import numpy as np
import matplotlib.pyplot as plt
# If you are running this on Notable, torch should be preinstalled
import torch
from torch import nn
# Install and import skorch, a sklearn-style interface to PyTorch
!pip install skorch
import skorch
from skorch import NeuralNetClassifier
sklearn is designed primarily for classical machine learning and does not provide end-to-end building blocks for neural network architectures such as convolutional networks (ConvNets) or Transformers. In practice, these models are implemented and trained using deep learning frameworks such as PyTorch or TensorFlow. To keep the tooling consistent, I won’t expect you to learn a full deep learning framework in the final taught week of the course.
That said, I will include an example using skorch (an sklearn-style interface to PyTorch) to show what a practical implementation looks like. You are not expected to be able to write or modify skorch/PyTorch code for this course.
10.1 Convolutional Neural Networks (ConvNets)¶
Last week, we looked at MLPs. These take in a vector input $\mathbf{x}\in\mathbb{R}^{D}$ and pass it through multiple layers. Each layer takes in some vector $\mathbf{h}^{(l-1)}\in\mathbb{R}^{H_{l-1}}$, multiplies it by a matrix, adds a bias, and then applies an elementwise nonlinearity:
$$\mathbf{h}^{(l)} = g(\mathbf{W}^{(l) }\mathbf{h}^{(l-1)} + \mathbf{b}^{(l)}) \,\,\text{for}\,\,l=1,2,\dots$$An image is represented as a 3D tensor $\mathbf{x}\in\mathbb{R}^{C\times H \times W}$. Vectorising an image for use in an MLP preserves the pixel values, but an MLP does not explicitly exploit their spatial arrangement. In the lecture, we looked at convolutional neural networks (ConvNets). These take in images as 3D tensors, and use a mixture of convolutional, pooling, and linear (fully connected) layers.
10.1.1 Convolutional layers¶
A convolutional layer is quite similar to a linear layer in an MLP, in that it consists of a linear transformation followed by a non-linearity. However, it takes in a 3D tensor $\mathbf{h}^{(l-1)}\in\mathbb{R}^{C_{l-1}\times H_{l-1}\times W_{l-1}}$ instead of a vector, and the linear transformation is a 2D convolution:
$$\mathbf{h}^{(l)} = g(\mathbf{W}^{(l) }\ast \mathbf{h}^{(l-1)} + \mathbf{b}^{(l)})\,\,\text{for}\,\,l=1,2,\dots$$Each $\mathbf{W}^{(l)}\in\mathbb{R}^{C_{l}\times C_{l-1}\times k\times k}$ consists of $C_{l}$ filters, each of dimensionality $C_{l-1}\times k \times k$ where $k$ is the kernel size.
If we define $\mathbf{z}^{(l)} = \mathbf{W}^{(l)}\ast \mathbf{h}^{(l-1)} \in\mathbb{R}^{C_{l}\times H_{l}\times W_{l}}$ then each of its elements $ \mathbf{z}^{(l)}_{c,\, i,\, j}$ is computed as
$$ \mathbf{z}^{(l)}_{c,\, i,\, j}=\sum_{c'=1}^{C_{l-1}} \sum_{u=0}^{k-1} \sum_{v=0}^{k-1} \mathbf{W}^{(l)}_{c,\, c',\, u,\, v}\; \mathbf{h}^{(l-1)}_{c',\, i+u,\, j+v}$$The bias $\mathbf{b}^{(l)}\in\mathbb{R}^{C_{l}}$ has one element per output channel. Elements of the preactivation $\mathbf{a}^{(l)} = \mathbf{z}^{(l)}+\mathbf{b}^{(l)}$ (where I have been rather lax with my maths, as I am broadcasting) are computed as
$$ \mathbf{a}^{(l)}_{c,\, i,\, j}= \mathbf{z}^{(l)}_{c,\, i,\, j} + b_{c}^{(l)}$$10.1.2 Pooling layers¶
Pooling layers reduce the spatial resolution of the input. Coupled with convolution layers, they can make ConvNets less sensitive to small translations; they do not guarantee complete translation invariance. For max pooling with kernel size and stride $k$, assuming $H$ and $W$ are divisible by $k$, if the input is $\mathbf{x}\in\mathbb{R}^{C \times H \times W}$ then the output is $\mathbf{y}\in\mathbb{R}^{C \times H/k \times W/k}$ where
$$ \mathbf{y}_{c,i,j}=\max_{\substack{0\le u < k\\[2pt] 0\le v < k}} \mathbf{x}_{c,\; i\,k + u,\; j\,k + v} $$For $k=2$ this means that each $2\times 2$ block in each channel is replaced by the maximum value across those 4 elements. It is also possible to perform average pooling. In this case each block is replaced by the mean value.
10.1.3 A small ConvNet for MNIST in skorch¶
Now that we have covered all the constituent parts of a ConvNet, I am going to provide code that trains the ConvNet below on MNIST using skorch and PyTorch. You will be implementing the same ConvNet in the lab, but using numpy instead.
First, let's load in MNIST, scale it to between 0 and 1, and reshape it so data points are 3D tensors (images) instead of vectors. Note that MNIST is greyscale so there is only 1 colour channel.
with open("data/mnist.npy", "rb") as f:
X = np.load(f, allow_pickle=True) # (70000, 784) uint8
y = np.load(f, allow_pickle=True).astype(np.int64) # (70000,)
# Scale to between 0 and 1 and change type for pytorch
X = (X/255.).astype(np.float32)
# Reshape to NCHW: (N, 1, 28, 28)
X = X.reshape(-1, 1, 28, 28)
X_train, y_train = X[:60000], y[:60000]
X_test, y_test = X[60000:], y[60000:]
Now let's create the ConvNet in PyTorch. There will be extra functionality for later.
# Attempt to fix RNG
torch.manual_seed(1)
np.random.seed(2)
class ConvNetModule(nn.Module):
def __init__(self, num_classes=10, in_shape=(1, 28, 28)):
super().__init__()
# Feature extractor (everything except the final linear classifier)
self.features = nn.Sequential(
nn.Conv2d(1, 8, kernel_size=5, padding=0),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(8, 16, kernel_size=5, padding=0),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
)
self.classifier = nn.Linear(256, num_classes)
def extract_features(self, x):
"""Return flattened feature vector (N, D) for feature extraction."""
x = self.features(x)
return torch.flatten(x, 1)
def forward(self, x):
z = self.extract_features(x)
return self.classifier(z)
Let's use skorch to train the ConvNet. Unfortunately, this will be slow on a CPU.
convnet = NeuralNetClassifier(
ConvNetModule,
max_epochs=10,
lr=0.002,
optimizer=torch.optim.Adam,
criterion=nn.CrossEntropyLoss,
)
convnet.fit(X_train, y_train)
Finally, we'll evaluate the accuracy of this ConvNet on the test set.
convnet.score(X_test, y_test)
ConvNets typically perform very well on MNIST because they exploit the spatial structure of images. Run the cell above to check the test accuracy of this model.
10.2 Feature extraction¶
A trained neural network can be used as a feature extractor. This is useful if we have a new, related task that can benefit from the information learnt in the original task, because the network has already learnt good features that can be reused instead of relearnt from scratch. For example, a vision model trained on a large image dataset will have learnt general visual features such as edges, textures, and shapes, which can be reused for a different image classification task. Similarly, a language model trained on large text corpora will have learnt useful linguistic features (e.g. word meaning and sentence structure) that can be reused for a new text classification task.
We can simply compute the features ($\phi({\mathbf{x}})$) for the data in our new task, and then apply e.g. logistic regression. In the lab, you will extract features using a network trained on 1 million images to separate out 1000 classes and see that they work rather well.
10.3 Fine-tuning¶
Fine-tuning goes a step beyond feature extraction by updating (refitting) some or all of the weights of a pretrained neural network so it adapts to a new task. In practice, this is usually done with a smaller learning rate than would be used for training from scratch, and it may involve freezing certain layers (i.e. preventing their parameters from being updated).
The key idea is that pretraining (often on a large, generic dataset) provides a strong starting point: the model has already learned useful representations, so fine-tuning can adapt these representations to a new task with far less data and compute than full training.
This is especially appealing because:
- It’s much cheaper (in time, compute, and energy) than training a new model from scratch
- It works well in low-data regimes, where you may not have enough labelled examples to train a deep network reliably
- It often improves performance compared with using the pretrained network as a fixed feature extractor, because the representations can shift toward what the new task needs
When fine-tuning, it is usually necessary to replace the "head" (the final task-specific layer). For example, a network trained on ImageNet typically ends with a classification head that maps the learned feature vector to 1000 logits (one per ImageNet class). If your new task instead requires, say, 200 classes, you would replace that final linear layer with a new one that outputs 200 logits, and then train that new head (alongside some or all of the previous layers) on your new dataset.
That's all folks!