MNIST neural network example

In the course sessions we trained a basic neural network to recognise handwritten numbers using Autograd and JAX. Here we continue investigating the JAX implementation.

Figure 1: Neural network schematic created with https://alexlenail.me/NN-SVG/index.html.

To begin with run the two blocks below to train the neural network. These do the same as in the second session of the course. The exercises from the second session are copied here afterwards. It would be a good idea to go back through them to remind yourself of the context of the MNIST problem.

%matplotlib inline

import matplotlib.image
import matplotlib.pyplot as plt
import numpy as np

# Functions and imports for training the handwriting classifier neural network with JAX

import array
import gzip
import os
import struct
from urllib.request import urlretrieve

import jax
import jax.numpy as jnp
from jax import grad, jit, random
from jax.scipy.special import logsumexp
import optax
import matplotlib.pyplot as plt

def load_mnist():
    """
    Download the MNIST training data and normalise it in preparation for training and testing.

    :return: number of images and arrays of normalised images and labels
    """
    print("Loading training data...")

    def download(url, filename):
        if not os.path.exists("data"):
            os.makedirs("data")
        out_file = os.path.join("data", filename)
        if not os.path.isfile(out_file):
            urlretrieve(url, out_file)

    def parse_labels(filename):
        with gzip.open(filename, "rb") as fh:
            magic, num_data = struct.unpack(">II", fh.read(8))
            return np.array(array.array("B", fh.read()), dtype=np.uint8)

    def parse_images(filename):
        with gzip.open(filename, "rb") as fh:
            magic, num_data, rows, cols = struct.unpack(">IIII", fh.read(16))
            return np.array(array.array("B", fh.read()), dtype=np.uint8).reshape(num_data, rows, cols)

    base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
    
    for filename in [
        "train-images-idx3-ubyte.gz",
        "train-labels-idx1-ubyte.gz",
        "t10k-images-idx3-ubyte.gz",
        "t10k-labels-idx1-ubyte.gz",
    ]:
        download(base_url + filename, filename)

    train_images = parse_images("data/train-images-idx3-ubyte.gz")
    train_labels = parse_labels("data/train-labels-idx1-ubyte.gz")
    test_images = parse_images("data/t10k-images-idx3-ubyte.gz")
    test_labels = parse_labels("data/t10k-labels-idx1-ubyte.gz")

    partial_flatten = lambda x: jnp.reshape(x, (x.shape[0], -1))
    one_hot = lambda x, k: jnp.array(x[:, None] == jnp.arange(k)[None, :], dtype=int)
    train_images = partial_flatten(train_images) / 255.0
    test_images = partial_flatten(test_images) / 255.0
    train_labels = one_hot(train_labels, 10)
    test_labels = one_hot(test_labels, 10)
    N_data = train_images.shape[0]
    return N_data, train_images, train_labels, test_images, test_labels

def init_random_params(scale, layer_sizes, key=random.PRNGKey(0)):
    """
    Build a list of (weights, biases) tuples, one for each layer in the net.

    :arg scale: scaling parameter
    :arg layer_sizes: sizes of each layer in the neural network
    :kwarg rs: initial random state
    :return: array of (weights, biases) tuples representing the net
    """
    keys = random.split(key, len(layer_sizes) - 1)
    params = []
    for k, m, n in zip(keys, layer_sizes[:-1], layer_sizes[1:]):
        W = scale * random.normal(k, (m, n))  # weight matrix
        b = scale * random.normal(k, (n))     # bias vector
        params.append((W, b))
    return params

def neural_net_predict(params, inputs):
    """
    Implements a deep neural network for classification.
    
    :arg params: list of (weights, bias) tuples
    :arg inputs: (N x D) input data matrix for images
    :return: class probabilities
    """
    for W, b in params[:-1]:
        inputs = jnp.tanh(inputs @ W + b)
    W, b = params[-1]
    outputs = inputs @ W + b
    return jnp.exp(outputs - logsumexp(outputs, axis=1, keepdims=True)) # softmax

def log_posterior(params, inputs, targets, L2_reg):
    """
    Compute a log-posterior objective function comprised of prior and likelihood components

    :arg params: list of (weights, bias) tuples
    :arg inputs: input data matrix for images
    :arg targets: target data for labels
    :arg L2_reg: scalar regularisation parameter
    :return: log-posterior objective function
    """
    l2_norm = sum(jnp.sum(W**2) + jnp.sum(b**2) for W, b in params)
    log_prior = -L2_reg * l2_norm
    probs = neural_net_predict(params, inputs)
    log_lik = jnp.sum(jnp.log(probs) * targets)
    return log_prior + log_lik

def accuracy(params, inputs, targets):
    """
    Compute the accuracy of a set of predictions.

    :arg params: list of (weights, bias) tuples
    :arg inputs: input data matrix for images
    :arg targets: target data for labels
    :return: mean accuracy of predictions
    """
    target_class = jnp.argmax(targets, axis=1)
    predicted_class = jnp.argmax(neural_net_predict(params, inputs), axis=1)
    return jnp.mean(predicted_class == target_class)

# Define training objective
def objective(params, inputs, targets, L2_reg=1.0):
    return -log_posterior(params, inputs, targets, L2_reg)

def print_perf(params, iteration):
    if iteration == 0:
        print("Epoch | Train accuracy | Test accuracy")
    if iteration % num_batches == 0:
        train_acc = accuracy(params, train_images, train_labels)
        test_acc = accuracy(params, test_images, test_labels)
        print(f"  {iteration // num_batches:1}   |  {train_acc:10.5f}    |  {test_acc:10.5f}")

def get_batch(iteration):
    idx = (iteration % num_batches) * batch_size
    x = train_images[idx:idx+batch_size]
    y = train_labels[idx:idx+batch_size]
    return x, y

def show_images(
    images,
    ims_per_row=5,
    padding=5,
    digit_dimensions=(28, 28),
    cmap="gray_r",
    vmin=0,
    vmax=None,
):
    """Images should be a (N_images x pixels) matrix."""
    N_images = images.shape[0]
    N_rows = (N_images - 1) // ims_per_row + 1
    pad_value = np.min(images.ravel())
    concat_images = np.full(
        (
            (digit_dimensions[0] + padding) * N_rows + padding,
            (digit_dimensions[1] + padding) * ims_per_row + padding,
        ),
        pad_value,
    )
    for i in range(N_images):
        cur_image = np.reshape(images[i, :], digit_dimensions)
        row_ix = i // ims_per_row
        col_ix = i % ims_per_row
        row_start = padding + (padding + digit_dimensions[0]) * row_ix
        col_start = padding + (padding + digit_dimensions[1]) * col_ix
        concat_images[
            row_start : row_start + digit_dimensions[0], col_start : col_start + digit_dimensions[1]
        ] = cur_image
    
    fig, axes = plt.subplots()
    cax = axes.matshow(concat_images, cmap=cmap, vmin=vmin, vmax=vmax)
    plt.xticks(np.array([]))
    plt.yticks(np.array([]))

# Model parameters
param_scale = 0.1
layer_sizes = [784, 200, 100, 10]
init_params = init_random_params(param_scale, layer_sizes)

# Training parameters
batch_size = 256
num_epochs = 5
step_size = 0.001

## Load training data
N, train_images, train_labels, test_images, test_labels = load_mnist()
num_batches = int(np.ceil(len(train_images) / batch_size))
print(f"Training dataset size: {N}")
print(f"Testing dataset size: {len(test_images)}")

## Train the neural network    
# Set up the solver
solver = optax.adam(step_size)
opt_state = solver.init(init_params)
params = init_params

# Get gradient of the objective function using JAX.
grad_fn = grad(objective)

@jit # Compile each step of the solver for better performance
def step(params, opt_state, inputs, targets):
    grads = grad_fn(params, inputs, targets)
    updates, opt_state = solver.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state

# Run the solver
for iteration in range(num_epochs * num_batches):
    inputs, targets = get_batch(iteration)
    params, opt_state = step(params, opt_state, inputs, targets)
    print_perf(params, iteration)
optimized_params = params

## Check the result
show_images(test_images[100:105])
predicted = neural_net_predict(optimized_params, test_images[100:105]).argmax(axis=1)
print(f"Predicted: {predicted}")

Effect of each pixel on the predicted result

These exercises are copied from session 2 of the course.

Starting from a blank image we can investigate how changing each pixel will influence the number predicted by the neural network.

As an aside, we can investigate what the neural network outputs for a blank image… although it is likely to be nonsense.

blank_image = jnp.zeros(28 * 28)

# The neural network expects and returns a 2D array of probabilities for each digit for multiple images
probs = neural_net_predict(optimized_params, blank_image[np.newaxis,:])[0]

# Show the probabilities for a blank image
fig = plt.figure(figsize=(5,3))
plt.bar(range(10), probs)
plt.xticks(range(10))
plt.ylabel("Digit probabilities")
plt.ylim([0, 1])
plt.show()

Exercise 1

What function, when differentiated, will show how each pixel influences the predicted digit? Fill this in in the code below.

Hint

The function should take the image as input and return the probabilities for each of the ten digits (the output of the neural network).

Solution

def f(x):
    """
    :arg x: input image
    :return: prediction of the digit this image represents
    """
    return neural_net_predict(optimized_params, x[np.newaxis, :])[0]

Question 2

We will be calculating the derivative of all outputs (digits) with respect to all inputs (pixels), i.e. the full Jacobian. Which differentiation approach will work best, forward mode or reverse mode? And which function should be used jax.grad, jax.jvp or jax.vjp?

Hint

Consider the number of inputs and outputs to the function. Generally, more inputs - reverse mode, more outputs - forward mode.

Solution

Reverse mode: Because there are more inputs (\(28 \times 28 = 784\) pixels) than outputs (\(10\)).

jax.vjp (vector-Jacobian product) because this is the driver for reverse mode. jax.jvp (Jacobian-vector product) is used for forward mode and the output is not a scalar so jax.grad cannot be used.

Exercise 3

Complete the code to get the gradients for each digit with respect to the pixels.

Hint

jax.vjp takes the function to differentiate and its input. It returns the function output (not needed) and a functon that calculates the vector-Jacobian product when given a seed vector.

The result of this is a list of gradients for each input parameter. There is only one in this case (\(x\)).

Solution

results = []
for digit, seed_vector in enumerate(seed_vectors):
    _, vjp_fn = jax.vjp(f, blank_image)
    result = vjp_fn(seed_vector)[0]
    results.append(result)

Or, more succinctly using jax.vmap:

_, vjp_fn = jax.vjp(f, blank_image)
results = jax.vmap(lambda v: vjp_fn(v)[0])(seed_vectors)

def f(x):
    """
    :arg x: input image
    :return: probabilities of each digit for this image
    """
    # Exercise 1: Define this
    return ???

# We will evaluate the function `f` at at array of zeros, which corresponds to a blank image
blank_image = jnp.zeros(28 * 28)

# We need seed vectors for each of the ten digits, we can use the rows of an identity matrix
seed_vectors = jnp.eye(10)

# Exercise 3: Get the effect of each pixel for each of the ten digits
results = []
for digit, seed_vector in enumerate(seed_vectors):
    result = ??? # TODO
    results.append(result)

# Plot the gradient fields
fig, axes = plt.subplots(1, 10, figsize=(14,2))
for digit, result in enumerate(results):
    vmax = np.max(np.abs(np.array(result)))
    ax = axes.ravel()[digit]
    ax.imshow(result.reshape(28, 28), cmap="RdBu_r", vmax=vmax, vmin=-vmax)
    ax.set_title(digit)
    ax.set_axis_off()

Effect of a perturbation

This exercise is copied from the second session of the course.

Next we will look at how making a perturbation to the image would impact the likelihood of each digit. Run the code below to make a perturbation.

# Choose an image
image = test_images[0]

# Function to differentiate
def f(x):
    return neural_net_predict(optimized_params, x[np.newaxis, :])[0]

# Define a perturbation
perturbation = 1.0 * jnp.zeros((28, 28)).at[:, 13:15].set(1.0).reshape(-1) # Vertical stripe
#perturbation = 0.2 * jnp.ones_like(image)                                 # Uniform
#perturbation = 0.2 * jax.random.normal(jax.random.PRNGKey(0), (28 * 28,)) # Random noise

# Plot
fig, axes = plt.subplots(1, 2, figsize=(6,3))

cmap = axes[0].imshow(image.reshape(28,28), cmap="gray_r")
plt.colorbar(cmap, ax=axes[0])
axes[0].set_title("Original image")
axes[0].axis("off")

cmap = axes[1].imshow((image + perturbation).reshape(28,28), cmap="gray_r")
plt.colorbar(cmap, ax=axes[1])
axes[1].set_title("After perturbation")
axes[1].axis("off")

plt.show()

Exercise 4

Complete the code below to use the gradient to estimate the change in the likelihood of each digit (output of the neural network) given the perturbation.

Hint

The gradient is calculated in a single direction in the input space (the perturbation). So the differentiation should use forward mode with jax.jvp.

The function f(x) only takes a single argument so the values and tangent passed to jax.jvp should each be a list of size 1.

Solution

probs, change_probs = jax.jvp(f, (image,), (perturbation,))

# Exercise 4: Estimate the change in likelihood of each digit as `change_probs`
probs, change_probs = ???

# Plot
fig, axes = plt.subplots(1, 2, figsize=(6,3))
axes[0].bar(range(10), probs)
axes[0].set_xticks(range(10))
axes[0].set_title("Original probabilities")
axes[0].semilogy()
axes[1].bar(range(10), change_probs / probs)
axes[1].set_xticks(range(10))
axes[1].set_title("Estimated relative change\nin probabilities")
plt.show()

Mini-project idea

We have considered image recognition but what about image generation? Pick a digit from 0-9 (say 9). Start with the blank image, an appropriately scaled random image, or any other grayscale image of appropriate dimensions. Can we use the gradient information to generate an image of your digit of choice?

What can we do to regularise the problem so that it generates something that actually looks like a digit? Regularisation will involve modifying the function that we are computing the derivative of. We want to penalise images that don’t look like digits.

Hint

Consider the average image in the training set. Can we modify the function to reduce the difference between the generated image and this?

mean_image = train_images.mean(axis=0)

Try out different initial images and see what the impact is on your generated image.

Can you think of any other things we could do? Can you get the confidence over 99%?

If you manage to achieve this, can you modify the approach to accept one of the digits 0-9 as input and return a generated image that the network would identify with confidence over 99%?

# Choose a digit
target = 9

def f(x):
    return neural_net_predict(optimized_params, x[np.newaxis, :])[0][target]  # TODO: Add regularisation term

# Choose a starting image
blank_image = jnp.zeros(28 * 28)
image = blank_image
# image = 0.2 * jax.random.normal(jax.random.PRNGKey(0), (28 * 28,))
# image = test_images[0]

# Repeatedly compute and add the gradient until confidence is sufficiently high
maxiter = 10000
atol = 1e-02
rtol = 1e-04
images = [image]
for i in range(maxiter):
    # Get the effect of each pixel
    y, vjp_fun = jax.vjp(f, image)
    dfdimage = vjp_fun(1.0)[0]
    image += dfdimage
    images.append(image)
    if 1.0 - atol < y <= 1.0:
        print(f"Converged in {i+1} iterations due to confidence exceeding tolerance")
        break
    if i == 0:
        image_ = image
    elif jnp.linalg.norm(image - image_) / jnp.linalg.norm(image_) < rtol:
        print(f"Image converged in {i+1} iterations")
        break
    image_ = image
else:
    raise RuntimeError(f"Failed to converge in {maxiter} iterations")

Plot the last few iterations

images = jnp.array(images)
show_images(images[-5:])
confidence = f(image)
print(f"Confidence that the digit is {target}: {100 * confidence:.2f}%")

Clearly, this looks nothing like a number 9! Let’s try to improve on this.