13-pytorch-03-model-def.qmd

---
title: PyTorch 03 - Model Definition
jupyter: python3
---

## Introduction

[![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/trgardos/ml-549-fa24/blob/main/13-pytorch-03-model-def.ipynb)

Based on PyTorch [Build Model Tutorial](https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html).


## Build the Neural Network

Neural networks comprise of layers/modules that perform operations on
data. 

The [torch.nn](https://pytorch.org/docs/stable/nn.html) namespace
provides all the building blocks you need to build your own neural
network. 

Every module in PyTorch subclasses the
[nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html).

A neural network is a module itself that consists of other modules
(layers). This nested structure allows for building and managing complex
architectures easily.

In the following sections, we'll build a neural network to classify
images like those in the **CIFAR-10** dataset.

```{python}
#| code-fold: true
import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
```

## Get Device for Training

We want to be able to train our model on a hardware accelerator like the
GPU or MPS, if available. 

Let's check to see if

* [torch.cuda](https://pytorch.org/docs/stable/notes/cuda.html) or
* [torch.backends.mps](https://pytorch.org/docs/stable/notes/mps.html) 

are available, otherwise we use the CPU.

```{python}
#| code-fold: false
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")
```

## Define the Class

We define our neural network by subclassing `nn.Module`, and initialize the
neural network layers in `__init__`. 

Every `nn.Module` subclass implements the operations on input data in the `forward` method.

```{python}
#| code-fold: false
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(3*32*32, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits
```

We create an instance of `NeuralNetwork`, and move it to the `device`,
and print its structure.

```{python}
#| code-fold: false
model = NeuralNetwork().to(device)
print(model)
```

::: {.callout-note}
Note that this is an _untrained_ model with randomly initialized weights.
:::

Keras has a nicer way to print the model architecture. To do something similar
in PyTorch, you can use `torchinfo` package.

```{python}
#| echo: false
import sys
if 'google.colab' in sys.modules:
    !pip install torchinfo
```

```{python}
#| code-fold: false
from torchinfo import summary
summary(model, input_size=(1, 3, 32, 32))
```

```{python}
# the torchinfo package seems to move the model to the CPU so reinitializing
model = NeuralNetwork().to(device)
```

To use the model, we pass it the input data. 

This executes the model's `forward`.

::: {.callout-caution}
Do not call `model.forward()` directly!
:::

```{python}
#| code-fold: false
X = torch.rand(1, 3, 32, 32, device=device)
logits = model(X)
# print(f"Logits: {logits.T}")
print(f"Logits size: {logits.size()}")
```

Calling the model on the input returns a $1\times10$ tensor
corresponding to each output of 10 raw predicted values for each class

::: {.callout-note}
Note that we had to set the batch size to 1 on the input, e.g. a $1\times3\times32\times32$ tensor.

That's because `nn.Flatten` is batch aware, i.e. it flattens each image in the batch individually.
:::

We get the prediction probabilities by passing it through an instance of the
`nn.Softmax` module.

```{python}
#| code-fold: false
pred_probab = nn.Softmax(dim=1)(logits)
# print(f"\nPredicted probabilities: {pred_probab.T}")
```

And then we get the predicted class by choosing the index with the maximum
probability.

```{python}
#| code-fold: false
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")
```

```{python}
#| echo: false
import sys
if 'google.colab' in sys.modules:
    !wget https://raw.githubusercontent.com/trgardos/ml-549-fa24/refs/heads/main/utils.py

```

#| echo: false
```{python}
#| echo: false
from utils import tensor_to_latex
from IPython.display import display, Math

latex_logits = tensor_to_latex(logits.T)
latex_pred_probab = tensor_to_latex(pred_probab.T)

display(Math(f"\\mathrm{{logits}} = {latex_logits}, \\quad \\mathrm{{pred\\_probab}} = {latex_pred_probab}"))
```


## Model Layers

Let's step through the layers in the model one by one.

To illustrate it, we will create a sample minibatch of 5 "images" with random
values of size 32x32 and see what happens to it as we pass it through the network.

```{python}
#| code-fold: false
input_image = torch.rand(5,3,32,32)
print(input_image.size())
```

### nn.Flatten

We initialize the
[nn.Flatten](https://pytorch.org/docs/stable/generated/torch.nn.Flatten.html)
layer to convert each 2D 3x32x32 image into a contiguous array of 3072
pixel values (the minibatch dimension (at dim=0) is maintained).

```{python}
#| code-fold: false
flatten = nn.Flatten()
flat_image = flatten(input_image)
print(flat_image.size())
```

### nn.Linear

The [linear
layer](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html)
is a module that applies a linear transformation on the input using its
stored weights and biases.

Remember from our Tensor discussions that the linear calculation for the linear
fully connected layer as in the following figure

![Deep Fully Connected Network](assets/images/pytorch/shallow_net.png)

can be expressed as a matrix multiplication.

$$
\begin{bmatrix}
w_{00} & w_{01} & w_{02} \\
w_{10} & w_{11} & w_{12} \\
w_{20} & w_{21} & w_{22} \\
w_{30} & w_{31} & w_{32} \\
\end{bmatrix}
\begin{bmatrix}
x_{0} \\
x_{1} \\
x_{2} \\
\end{bmatrix} +
\begin{bmatrix}
b_{0} \\
b_{1} \\
b_{2} \\
b_{3} \\
\end{bmatrix}
=
\begin{bmatrix}
h_{0} \\
h_{1} \\
h_{2} \\
h_{3} \\
\end{bmatrix}
$$

In the example below the weights matrix is $20\times3072$ and the bias vector is $20\times1$.

```{python}
#| collapsed: false
layer1 = nn.Linear(in_features=3*32*32, out_features=20)
hidden1 = layer1(flat_image)
print(hidden1.size())
```

We get a 5 batches of 20 values each in a $5\times20$ tensor.

### nn.ReLU

Non-linear activations are what create the complex mappings between the
model's inputs and outputs. 

They are applied after linear transformations to introduce *nonlinearity*,
helping neural networks learn a wide variety of phenomena.

Without non-linear activations, the model would just be a linear
function.

Later we'll see that adding non-linearities make neural networks **"Universal Approximators"**.

In this model, we use
[nn.ReLU](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html)
between our linear layers, but there are other activations to introduce
non-linearity in your model.

The ReLU activation function is defined as $ReLU(x) = \max(0, x)$.

```{python}
#| code-fold: true
import matplotlib.pyplot as plt
import numpy as np

# Generate values between -5 and 5
x = np.linspace(-5, 5, 100)

# Apply ReLU function
tensor_x = torch.tensor(x)
tensor_y = nn.ReLU()(tensor_x)
y = tensor_y.numpy()

# or equivalently
# y = nn.ReLU()(torch.tensor(x)).numpy()

# Plot the ReLU function
plt.figure(figsize=(8, 6))
plt.plot(x, y, label="ReLU(x)")
plt.title("ReLU Activation Function")
plt.xlabel("Input")
plt.ylabel("Output")
plt.legend()
plt.grid(True)
plt.show()
```

```{python}
#| code-fold: false
print(f"Before ReLU: {hidden1}\n\n")
hidden1 = nn.ReLU()(hidden1)
print(f"After ReLU: {hidden1}")
```

### nn.Sequential

Although not technically a layer, 
[nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html)
is an ordered container of modules and can be treated as a layer.

The data is passed through all the modules in the same order as defined.
You can use sequential containers to put together a quick network like
`seq_modules`.

```{python}
#| code-fold: false
seq_modules = nn.Sequential(
    flatten,
    layer1,
    nn.ReLU(),
    nn.Linear(20, 10)
)
input_image = torch.rand(5,3,32,32)
logits = seq_modules(input_image)
```

### nn.Softmax

The last linear layer of the neural network returns _logits_
- raw values in $(-\infty, \infty)$ - which are passed to the
[nn.Softmax](https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html)
module. 

The softmax function is defined as

$$
\mathrm{Softmax}(x_{i}) = \frac{e^{x_i}}{\sum_j e^{x_j}}
$$

The logits are scaled to values $[0, 1]$ representing the
model's predicted probabilities for each class. `dim` parameter
indicates the dimension along which the values must sum to 1.

```{python}
#| code-fold: false
softmax = nn.Softmax(dim=1)
pred_probab = softmax(logits)
print(f"Softmax size: {pred_probab.size()}")
```

## Model Parameters

Many layers inside a neural network are *parameterized*, i.e. have
associated weights and biases that are optimized during training.

Subclassing `nn.Module` automatically tracks all fields defined inside
your model object, and makes all parameters accessible using your
model's `parameters()` or `named_parameters()` methods.

In this example, we iterate over each parameter, and print its size and
a preview of its values.

```{python}
#| code-fold: false
print(f"Model structure: {model}\n\n")

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()}") # | Values : {param[:2]} \n")
```

We can count the sizes of the matrices and vectors to verify that the
number of parameters in the model agrees with the summary printed earlier.

```{python}
print(f"{512*3072 + 512 + 512*512 + 512 + 512*10 + 10:,}")
```

## CIFAR-10 Images

Let's load the CIFAR-10 training data so we can apply our model to it.

```{python}
import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

training_data = datasets.CIFAR10(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)
```


We can then select a random sample from the training data and print its shape,
type, and label.

```{python}
sample_idx = torch.randint(len(training_data), size=(1,)).item()
img, label = training_data[sample_idx]

print(f"Sample index: {sample_idx}")
print(f"Label: {label}")
print(f"img.shape: {img.shape}")
print(f"img.dtype: {img.dtype}")
```

We can then create a `DataLoader` to load the data in batches.

```{python}
from torch.utils.data import DataLoader
train_dataloader = DataLoader(training_data, batch_size=5, shuffle=True)
```

Let's grab one batch of images and labels and print their shapes.

```{python}
# Display image and label.
train_features, train_labels = next(iter(train_dataloader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")
```

```{python}
labels_map = {
    0: "plane",
    1: "car",
    2: "bird",
    3: "cat",
    4: "deer",
    5: "dog",
    6: "frog",
    7: "horse",
    8: "ship",
    9: "truck",
}
import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 5, figsize=(15, 3))
for i in range(5):
    img = train_features[i].squeeze() # squeeze() removes dimension of size 1, e.g. (1, 3, 32, 32) -> (3, 32, 32)
    label = train_labels[i]

    axes[i].imshow(img.permute(1, 2, 0))
    axes[i].set_title(labels_map[label.item()])
    axes[i].axis('off')

plt.show()

```

Let's instantiate our model again and move it to the device.

```{python}
#| code-fold: false
model = NeuralNetwork().to(device)
```

We can then pass the batch of images through the model and print the size of the
logits.

::: {.callout-important}
Note that we had to move the tensor representing thebatch of images to the device.
:::

```{python}
#| code-fold: false
logits = model(train_features.to(device))
print(f"Logits size: {logits.size()}")
```

And finally, we can pass the logits through the softmax function to get the
predicted probabilities.

```{python}
#| code-fold: false
softmax = nn.Softmax(dim=1)
pred_probab = softmax(logits)
print(f"Softmax size: {pred_probab.size()}")
```


```{python}
#| echo: false
from utils import tensor_to_latex
from IPython.display import display, Math

latex_logits = tensor_to_latex(logits.T)
latex_pred_probab = tensor_to_latex(pred_probab.T)

display(Math(f"\\mathrm{{logits}} = {latex_logits}, \\quad \\mathrm{{pred\\_probab}} = {latex_pred_probab}"))
```

And like before we get the predicted class by choosing the index with the maximum
probability.

```{python}
#| code-fold: false
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")
```

And of course we are not getting the right predictions yet because we have not
trained the model!

Of course, in practice we would use other types of layers, e.g. convolutions,
dropout, batch normalization, etc., especially for image data.

## Further Reading

-   [torch.nn API](https://pytorch.org/docs/stable/nn.html)