Neural Network with JAX

# HANDSON: We would like to do some comparison ...
#          Let's separate the class into three groups.
#          Group 1 will use CPU for running this notebook.
#          Groups 2 and 3 will use GPU and TPU, respectively.

from os.path import isfile
from urllib.request import urlretrieve

base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"

# File names
files = {
    "train_images": "train-images-idx3-ubyte.gz",
    "train_labels": "train-labels-idx1-ubyte.gz",
    "test_images":  "t10k-images-idx3-ubyte.gz",
    "test_labels":  "t10k-labels-idx1-ubyte.gz",
}

for key, file in files.items():
    if not isfile(file):
        url = base_url + file
        print(f"Downloading {url} to {file}...")
        urlretrieve(url, file)
    else:
        print(f"{file} exists; skip download")

import gzip
import struct
import array
from jax import numpy as np

# Parsing functions

def parse_labels(file):
    with gzip.open(file, "rb") as fh:
        _magic, num_data = struct.unpack(">II", fh.read(8))
        # Read the label data as 1-byte unsigned integers
        return np.array(array.array("B", fh.read()), dtype=np.uint8)

def parse_images(file):
    with gzip.open(file, "rb") as fh:
        _magic, num_data, rows, cols = struct.unpack(">IIII", fh.read(16))
        # Read the image data as 1-byte unsigned integers
        images = np.array(array.array("B", fh.read()), dtype=np.uint8)
        # Reshape to (num_data, 28, 28)
        images = images.reshape(num_data, rows, cols)
        return images

# Parse raw data

train_images_raw = parse_images(files["train_images"])
train_labels_raw = parse_labels(files["train_labels"])

test_images_raw  = parse_images(files["test_images"])
test_labels_raw  = parse_labels(files["test_labels"])

# Standardize the images, i.e., flatten and normalize images to [0, 1]
def standardize(images):
    return images.reshape(-1, 28*28).astype(np.float32) / 255

train_images = standardize(train_images_raw)
test_images  = standardize(test_images_raw)

# One-hot encode labels
def one_hot(labels, num_classes=10):
    return np.eye(num_classes)[labels]

train_labels = one_hot(train_labels_raw, 10).astype(np.float32)
test_labels  = one_hot(test_labels_raw,  10).astype(np.float32)

from matplotlib import pyplot as plt

plt.imshow(train_images_raw[0,:,:], cmap='gray')

plt.imshow(test_images_raw[0,:,:], cmap='gray')

from numpy.random import RandomState

def init_params(scale, layer_sizes, rng=RandomState(0)):
    """
    Initialize the parameters (weights and biases) for each layer in the network.

    Parameters
    ----------
    scale : float
        A scaling factor to control the initial range of the weights.
    layer_sizes : list of int
        The sizes of each layer in the network.
        e.g., [784, 1024, 1024, 10] means:
            - Input layer: 784 units
            - Hidden layer 1: 1024 units
            - Hidden layer 2: 1024 units
            - Output layer: 10 units
    rng : numpy.random.RandomState
        Random state for reproducibility.

    Returns
    -------
    params : list of tuples (W, b)
        Each tuple contains (W, b) for a layer.
        - W is a (input_dim, output_dim) array of weights
        - b is a (output_dim,) array of biases
    """
    return [
        (scale * rng.randn(m, n), scale * rng.randn(n))
        for m, n in zip(layer_sizes[:-1], layer_sizes[1:])
    ]

# Define network architecture and hyperparameters

layer_sizes = [784, 1024, 1024, 10]  # 2 hidden layers
param_scale = 0.1

# Initialize parameters

params = init_params(param_scale, layer_sizes)

# HANDSON: count how many parameters are in your neural network!

# HANDSON: Let's split the class into 2 groups.
#          Group 1 makes their neural networks wider.
#          Group 2 makes theirs deeper.
#          We will compare the training and testing performance later.

from jax import numpy as np
from jax.scipy.special import logsumexp

def predict(params, inputs):
    """
    Compute the network's output logits for a batch of inputs, then subtract
    log-sum-exp for numerical stability (log-softmax).

    Network architecture:
      - Hidden layers use tanh activation
      - Output layer is linear (we'll do log-softmax here)

    Parameters
    ----------
    params : list of (W, b) tuples
        Network's parameters for each layer.
    inputs : np.ndarray
        A batch of input data of shape (batch_size, input_dim).

    Returns
    -------
    np.ndarray
        Log probabilities of shape (batch_size, 10).
    """
    activations = inputs

    # Hidden layers
    for w, b in params[:-1]:
        outputs = np.dot(activations, w) + b
        activations = np.tanh(outputs)

    # Final layer (logits)
    final_w, final_b = params[-1]
    logits = np.dot(activations, final_w) + final_b

    # Log-Softmax: subtract logsumexp for numerical stability
    return logits - logsumexp(logits, axis=1, keepdims=True)

def loss(params, batch):
    """
    Computes the average negative log-likelihood loss for a batch.

    Parameters
    ----------
    params : list of (W, b) tuples
        The network's parameters.
    batch : tuple (inputs, targets)
        - inputs: np.ndarray of shape (batch_size, 784)
        - targets: np.ndarray of shape (batch_size, 10) (one-hot labels)

    Returns
    -------
    float
        Mean negative log-likelihood over the batch.
    """
    inputs, targets = batch
    preds = predict(params, inputs)
    
    # preds are log-probs, multiply with one-hot targets and sum -> log-likelihood
    return -np.mean(np.sum(preds * targets, axis=1))

def accuracy(params, batch):
    """
    Computes classification accuracy of the network on a given batch.

    Parameters
    ----------
    params : list of (W, b) tuples
        The network's parameters.
    batch : tuple (inputs, targets)
        - inputs: np.ndarray (batch_size, 784)
        - targets: np.ndarray (batch_size, 10) (one-hot labels)

    Returns
    -------
    float
        Fraction of correctly classified samples.
    """
    inputs, targets = batch
    target_class = np.argmax(targets, axis=1)  # ground truth index
    predicted_class = np.argmax(predict(params, inputs), axis=1)
    return np.mean(predicted_class == target_class)

from jax import jit, grad

@jit
def update(params, batch, step_size):
    """
    Single step of gradient-based parameter update using simple SGD.

    grad(loss)(params, batch) computes the gradient of the loss function
    with respect to the parameters for the given batch.
    """
    grads = grad(loss)(params, batch)
    return [
        (w - step_size * dw, b - step_size * db)
        for (w, b), (dw, db) in zip(params, grads)
    ]

batch_size = 128  # the number of samples per parameter update.

num_train   = train_images.shape[0]
num_batches = num_train // batch_size

def get_batch(rng=RandomState(0)):
    """
    Generator function that yields shuffled batches indefinitely.
    """
    while True:
        # Randomly permute the indices
        perm = rng.permutation(num_train)
        for i in range(num_batches):
            batch_idx = perm[i*batch_size : (i+1)*batch_size]
            # Yield a tuple (inputs, labels) for this batch
            yield (train_images[batch_idx], train_labels[batch_idx])

train_batch_generator = get_batch()

learning_rate = 0.001  # the number of times we iterate over the entire training dataset.
num_epochs    = 5      # the number of samples per parameter update.

from time import time

for epoch in range(num_epochs):
    start_time = time()

    # Go through the entire training set
    for _ in range(num_batches):
        batch_data = next(train_batch_generator)
        params     = update(params, batch_data, step_size=learning_rate)

    epoch_time = time() - start_time

    # Evaluate training and test accuracy
    train_acc = accuracy(params, (train_images, train_labels))
    test_acc  = accuracy(params, (test_images, test_labels))

    print(f"Epoch {epoch} in {epoch_time:0.2f} sec")
    print(f"Training accuracy: {train_acc:.4f}")
    print(f"Test accuracy:     {test_acc:.4f}")

# HANDSON: At this point, we have a fully operational training
#          pipeline for MNIST.
#          You can experiment with different hyperparameters (e.g.,
#          learning rate, batch size, number of epochs, network
#          architecture) to see how they affect model performance.

from PIL import Image

def load_image(file):
    """
    Loads a grayscale image from `file`, resizes it to 28x28,
    and converts it to a (784,) float32 array with values in [0, 1].

    Parameters
    ----------
    file : str
        Path to the image file (e.g., a PNG or JPG).

    Returns
    -------
    np.ndarray
        An array of shape (1, 784) containing normalized pixel values
        suitable as input to our trained model.
    """
    # Convert the image to grayscale and resize to 28x28
    img = Image.open(file).convert('L').resize((28, 28))

    # Convert to a NumPy array and normalize pixel intensities to [0, 1]
    arr = np.array(img).astype(np.float32) / 255.0

    # Flatten the 28x28 image into a single 784-dimensional vector
    arr = arr.flatten()

    # Reshape to (1, 784) to match the model's expected input batch shape
    return np.array([arr])

def predict_digit(params, file_path):
    """
    Predict the digit class for a custom handwritten image.

    Parameters
    ----------
    params : list of (W, b) tuples
        The trained network parameters.
    file_path : str
        Path to the custom image file.

    Returns
    -------
    int
        The predicted digit label (0 through 9).
    """
    # Convert the image to a suitable NumPy array
    arr_np = load_image(file_path)
    
    # Use our 'predict' function to get log probabilities for each class
    log_probs = predict(params, arr_np)  # shape: (1, 10)
    
    # Select the digit class with the highest log probability
    return int(np.argmax(log_probs, axis=1)[0])

# Download the sample image from GitHub

! wget https://raw.githubusercontent.com/ua-2025q3-astr501-513/ua-2025q3-astr501-513.github.io/refs/heads/main/501/04/sample.png

# Predict what the digit is

label = predict_digit(params, "sample.png")
print(f"Predicted digit: {label}")

# HANDSON: 1. Capture or scan your handwritten digit and save as
#             "sample.png".
#             You may also use your cell phone for this.
#
#          2. Call the function:
#
#                 label = predict_digit(params, "sample.png")
#                 print(f"Predicted digit: {label}")
#
#          3. Inspect the result:
#             See whether the predicted label matches the actual digit
#             you wrote.
#
#             With these two functions, your MNIST-trained model can
#             be used in real-world testing scenarios, allowing you to
#             evaluate its performance on custom, hand-drawn images.