---
title: "pytorch - optim & nn"
subtitle: "Lecture 23"
author: "Dr. Colin Rundel"
footer: "Sta 663 - Spring 2023"
format:
  revealjs:
    theme: slides.scss
    transition: fade
    slide-number: true
    self-contained: true
execute: 
  echo: true
---

```{python}
#| label: setup
#| include: false
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import scipy

import statsmodels.formula.api as smf

import torch

import os
import math

plt.rcParams['figure.dpi'] = 200

np.set_printoptions(
  edgeitems=30, linewidth=60,
  precision = 5, suppress=True
)

pd.set_option("display.width", 130)
pd.set_option("display.max_columns", 10)
pd.set_option("display.precision", 6)
```

```{r}
#| label: r_setup
#| include: false
knitr::opts_chunk$set(
  fig.align="center",
  cache=FALSE
)

local({
  hook_err_old <- knitr::knit_hooks$get("error")  # save the old hook
  knitr::knit_hooks$set(error = function(x, options) {
    # now do whatever you want to do with x, and pass
    # the new x to the old hook
    x = sub("## \n## Detailed traceback:\n.*$", "", x)
    x = sub("Error in py_call_impl\\(.*?\\)\\: ", "", x)
    #x = stringr::str_wrap(x, width = 100)
    hook_err_old(x, options)
  })
  
  hook_warn_old <- knitr::knit_hooks$get("warning")  # save the old hook
  knitr::knit_hooks$set(warning = function(x, options) {
    x = sub("<string>:1: ", "", x)
    #x = stringr::str_wrap(x, width = 100)
    hook_warn_old(x, options)
  })
  
  hook_msg_old <- knitr::knit_hooks$get("output")  # save the old hook
  knitr::knit_hooks$set(output = function(x, options) {
    x = stringr::str_replace(x, "(## ).* ([A-Za-z]+Warning:)", "\\1\\2")
    x = stringr::str_split(x, "\n")[[1]]
    #x = stringr::str_wrap(x, width = 120, exdent = 3)
    x = stringr::str_remove_all(x, "\r")
    x = stringi::stri_wrap(x, width=120, exdent = 3, normalize=FALSE)
    x = paste(x, collapse="\n")
    
    #x = stringr::str_wrap(x, width = 100)
    hook_msg_old(x, options)
  })
})
```

# Odds & Ends

## Torch models

:::: {.columns .small}
::: {.column width='50%'}
Implementation details:

* Models are implemented as a class - inherits from `torch.nn.Module`

* Must implement constructor and `forward()` method

  * `__init__()` should call parent constructor via `super()`
    
    * Use `torch.nn.Parameter()` to indicate model parameters
  
  * `forward()` should implement the model - constants + parameters -> predictions

:::

::: {.column width='50%' .fragment}
Fitting proceedure:

* For each iteration of solver:

  * Get current predictions via a call to `forward()` or equivalent.
  
  * Calculate a loss or equivalent (scalar)

  * Call `backward()` method on loss
  
  * Use built-in optimizer (`step()` and `zero_grad()`)
:::
::::


## From last time

::: {.small}
```{python}
#| code-line-numbers: "|18"
class Model(torch.nn.Module):
    def __init__(self, X, y, beta=None):
        super().__init__()
        self.X = X
        self.y = y
        if beta is None:
          beta = torch.zeros(X.shape[1])
        beta.requires_grad = True
        self.beta = torch.nn.Parameter(beta)
        
    def forward(self, X):
        return X @ self.beta
    
    def fit(self, opt, n=1000, loss_fn = torch.nn.MSELoss()):
      losses = []
      for i in range(n):
          loss = loss_fn(
            self(self.X).squeeze(), 
            self.y.squeeze()
          )
          loss.backward()
          opt.step()
          opt.zero_grad()
          losses.append(loss.item())
      
      return losses
```
:::


## What is `self(self.X)`?

This is (mostly) just short hand for calling `self.forward(X)` to generate the output tensors from the current value(s) of the parameters. This is done via the special `__call__()` method in the `torch.nn.Module` class. `__call__()` allows python classes to be invoked like functions.

. . .

<br/>

:::: {.columns .small}
::: {.column width='50%'}
```{python}
class greet:
  def __init__(self, greeting):
    self.greeting = greeting
  def __call__(self, name):
    return self.greeting + " " + name
```
:::

::: {.column width='50%'}
```{python}
hello = greet("Hello")
hello("Jane")

gm = greet("Good morning")
gm("Bob")
```
:::
::::

::: {.aside}
Using a class like a function is often refered to as a functor in other languages
:::



# MNIST & Logistic models

## MNIST handwritten digits - simplified

::: {.small}
```{python}
from sklearn.datasets import load_digits
digits = load_digits()
```
:::


:::: {.columns .small}
::: {.column width='50%'}
```{python}
X = digits.data
X.shape
X[0:2]
```
:::

::: {.column width='50%'}
```{python}
y = digits.target
y.shape
y[0:10]
```
:::
::::


## Example digits

```{python}
#| echo: false
#| out.width: 85%
fig, axes = plt.subplots(nrows=5, ncols=10, figsize=(10, 6), layout="constrained")
axes2 = [ax for row in axes for ax in row]

for ax, image, label in zip(axes2, digits.images, digits.target):
    ax.set_axis_off()
    img = ax.imshow(image, cmap=plt.cm.gray_r, interpolation="nearest")
    txt = ax.set_title(f"{label}")
    
plt.show()
```

## Test train split

::: {.small}
```{python}
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.20, shuffle=True, random_state=1234
)
```
:::

:::: {.columns .small}
::: {.column width='50%'}
```{python}
X_train.shape
y_train.shape
X_test.shape
y_test.shape
```
:::

::: {.column width='50%' .fragment}
```{python}
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
```
```{python}
lr = LogisticRegression(
  penalty=None
).fit(
  X_train, y_train
)
```
```{python}
accuracy_score(y_train, lr.predict(X_train))
accuracy_score(y_test, lr.predict(X_test))
```
:::
::::



## As Tensors

::: {.small}
```{python}
X_train = torch.from_numpy(X_train).float()
y_train = torch.from_numpy(y_train)
X_test = torch.from_numpy(X_test).float()
y_test = torch.from_numpy(y_test)
```

```{python}
X_train.shape
y_train.shape
X_test.shape
y_test.shape
```
:::


## PyTorch Model

::: {.small}
```{python}
#| code-line-numbers: "|2-10|11-12|14-24"
class mnist_model(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.beta = torch.nn.Parameter(
          torch.randn(input_dim, output_dim, requires_grad=True)  
        )
        self.intercept = torch.nn.Parameter(
          torch.randn(output_dim, requires_grad=True)  
        )
        
    def forward(self, X):
        return (X @ self.beta + self.intercept).squeeze()
    
    def fit(self, X_train, y_train, X_test, y_test, lr=0.001, n=1000):
      opt = torch.optim.SGD(self.parameters(), lr=lr, momentum=0.9) 
      losses = []
      
      for i in range(n):
          opt.zero_grad()
          loss = torch.nn.CrossEntropyLoss()(self(X_train), y_train)
          loss.backward()
          opt.step()
          
          losses.append(loss.item())
      
      return losses
```
:::

## Cross entropy loss

::: {.small}
```{python}
model = mnist_model(64, 10)
l = model.fit(X_train, y_train, X_test, y_test)
```
:::

```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
plt.plot(l, label="loss")
plt.legend()
plt.show()
```

## Out-of-sample accuracy

:::: {.columns .small}
::: {.column width='50%'}
```{python}
model(X_test)
```
:::

::: {.column width='50%' .fragment}
```{python}
val, index = torch.max(model(X_test), dim=1)
index
```
:::
::::

::: {.small}
```{python}
(index == y_test).sum()
(index == y_test).sum() / len(y_test)
```
:::



## Calculating Accuracy

::: {.small}
```{python}
#| code-line-numbers: "|25-30"
class mnist_model(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.beta = torch.nn.Parameter(
          torch.randn(input_dim, output_dim, requires_grad=True)  
        )
        self.intercept = torch.nn.Parameter(
          torch.randn(output_dim, requires_grad=True)  
        )
        
    def forward(self, X):
        return (X @ self.beta + self.intercept).squeeze()
    
    def fit(self, X_train, y_train, X_test, y_test, lr=0.001, n=1000, acc_step=10):
      opt = torch.optim.SGD(self.parameters(), lr=lr, momentum=0.9) 
      losses, train_acc, test_acc = [], [], []
      
      for i in range(n):
          opt.zero_grad()
          loss = torch.nn.CrossEntropyLoss()(self(X_train), y_train)
          loss.backward()
          opt.step()
          losses.append(loss.item())
          
          if (i+1) % acc_step == 0:
            val, train_pred = torch.max(self(X_train), dim=1)
            val, test_pred  = torch.max(self(X_test), dim=1)
            
            train_acc.append( (train_pred == y_train).sum() / len(y_train) )
            test_acc.append( (test_pred == y_test).sum() / len(y_test) )
            
      return (losses, train_acc, test_acc)
```
:::

## Performance

::: {.small}
```{python}
loss, train_acc, test_acc = mnist_model(
  64, 10
).fit(
  X_train, y_train, X_test, y_test,acc_step=10, n=3000
)
```
:::


```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
plt.plot(range(0,3000,10), train_acc, label="train accuracy")
plt.plot(range(0,3000,10), test_acc, label="test accuracy")
plt.legend()
plt.show()
```


## NN Layers

::: {.small}
```{python}
#| code-line-numbers: "|4,7"
class mnist_nn_model(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = torch.nn.Linear(input_dim, output_dim)
        
    def forward(self, X):
        return self.linear(X)
    
    def fit(self, X_train, y_train, X_test, y_test, lr=0.001, n=1000, acc_step=10):
      opt = torch.optim.SGD(self.parameters(), lr=lr, momentum=0.9) 
      losses, train_acc, test_acc = [], [], []
      
      for i in range(n):
          opt.zero_grad()
          loss = torch.nn.CrossEntropyLoss()(self(X_train), y_train)
          loss.backward()
          opt.step()
          losses.append(loss.item())
          
          if (i+1) % acc_step == 0:
            val, train_pred = torch.max(self(X_train), dim=1)
            val, test_pred  = torch.max(self(X_test), dim=1)
            
            train_acc.append( (train_pred == y_train).sum() / len(y_train) )
            test_acc.append( (test_pred == y_test).sum() / len(y_test) )
            
      return (losses, train_acc, test_acc)
```
:::


## NN linear layer


Applies a linear transform to the incoming data ($x$):
$$y = x A^T+b$$

. . .

::: {.small}
```{python}
X.shape
model = mnist_nn_model(64, 10)
model.parameters()
list(model.parameters())[0].shape  # A - weights (betas)
list(model.parameters())[1].shape  # b - bias
```
:::



## Performance

::: {.small}
```{python}
loss, train_acc, test_acc = model.fit(X_train, y_train, X_test, y_test, n=1500)
```
:::

:::: {.columns .small}
::: {.column width='50%'}
```{python}
train_acc[-5:]
```
:::

::: {.column width='50%'}
```{python}
test_acc[-5:]
```
:::
::::



```{python}
#| echo: false
plt.figure(figsize=(12,6))
plt.subplot(121)
plt.plot(loss, label="loss")
plt.legend()

plt.subplot(122)
plt.plot(range(0,1500,10), train_acc, label="train accuracy")
plt.plot(range(0,1500,10), test_acc, label="test accuracy")
plt.legend()

plt.show()
```


# Feedforward Neural Network

## FNN Model

::: {.small}
```{python}
#| code-line-numbers: "|4-6,9-12"
class mnist_fnn_model(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, nl_step = torch.nn.ReLU(), seed=1234):
        super().__init__()
        self.l1 = torch.nn.Linear(input_dim, hidden_dim)
        self.nl = nl_step
        self.l2 = torch.nn.Linear(hidden_dim, output_dim)
        
    def forward(self, X):
        out = self.l1(X)
        out = self.nl(out)
        out = self.l2(out)
        return out
    
    def fit(self, X_train, y_train, X_test, y_test, lr=0.001, n=1000, acc_step=10):
      opt = torch.optim.SGD(self.parameters(), lr=lr, momentum=0.9) 
      losses, train_acc, test_acc = [], [], []
      
      for i in range(n):
          opt.zero_grad()
          loss = torch.nn.CrossEntropyLoss()(self(X_train), y_train)
          loss.backward()
          opt.step()
          
          losses.append(loss.item())
          
          if (i+1) % acc_step == 0:
            val, train_pred = torch.max(self(X_train), dim=1)
            val, test_pred  = torch.max(self(X_test), dim=1)
            
            train_acc.append( (train_pred == y_train).sum() / len(y_train) )
            test_acc.append( (test_pred == y_test).sum() / len(y_test) )
            
      return (losses, train_acc, test_acc)
```
:::

## Non-linear activation functions

:::: {.columns .small}
::: {.column width='50%'}
$$\text{Tanh}(x) = \frac{\exp(x)-\exp(-x)}{\exp(x) + \exp(-x)}$$
:::

::: {.column width='50%'}
$$\text{ReLU}(x) = \max(0,x)$$
:::
::::

:::: {.columns .small}
::: {.column width='50%'}
![](imgs/torch_Tanh.png){fig-align="center" width="100%"}
:::

::: {.column width='50%'}
![](imgs/torch_ReLU.png){fig-align="center" width="100%"}
:::
::::


## Model parameters

```{python}
model = mnist_fnn_model(64,64,10)
len(list(model.parameters()))
for i, p in enumerate(model.parameters()):
  print("Param", i, p.shape)
```

## Performance - ReLU

::: {.small}
```{python}
loss, train_acc, test_acc = mnist_fnn_model(64,64,10).fit(
  X_train, y_train, X_test, y_test, n=2000
)
train_acc[-5:]
test_acc[-5:]
```
:::

```{python}
#| echo: false
#| out-width: 75%
plt.figure(figsize=(12,6))
plt.subplot(121)
plt.plot(loss, label="loss")
plt.legend()

plt.subplot(122)
plt.plot(range(0,2000,10), train_acc, label="train accuracy")
plt.plot(range(0,2000,10), test_acc, label="test accuracy")
plt.legend()

plt.show()
```



## Performance - tanh

::: {.small}
```{python}
loss, train_acc, test_acc = mnist_fnn_model(
  64,64,10, nl_step=torch.nn.Tanh()
).fit(
  X_train, y_train, X_test, y_test, n=2000
)
train_acc[-5:]
test_acc[-5:]
```
:::

```{python}
#| echo: false
#| out.width: 75%
plt.figure(figsize=(12,6))
plt.subplot(121)
plt.plot(loss, label="loss")
plt.legend()

plt.subplot(122)
plt.plot(range(0,2000,10), train_acc, label="train accuracy")
plt.plot(range(0,2000,10), test_acc, label="test accuracy")
plt.legend()

plt.show()
```


## Adding another layer

::: {.small}
```{python}
#| code-line-numbers: "|4-8,11-16"
class mnist_fnn2_model(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, nl_step = torch.nn.ReLU(), seed=1234):
        super().__init__()
        self.l1 = torch.nn.Linear(input_dim, hidden_dim)
        self.nl = nl_step
        self.l2 = torch.nn.Linear(hidden_dim, hidden_dim)
        self.nl = nl_step
        self.l3 = torch.nn.Linear(hidden_dim, output_dim)
        
    def forward(self, X):
        out = self.l1(X)
        out = self.nl(out)
        out = self.l2(out)
        out = self.nl(out)
        out = self.l3(out)
        return out
    
    def fit(self, X_train, y_train, X_test, y_test, lr=0.001, n=1000, acc_step=10):
      loss_fn = torch.nn.CrossEntropyLoss()
      opt = torch.optim.SGD(self.parameters(), lr=lr, momentum=0.9) 
      losses, train_acc, test_acc = [], [], []
      
      for i in range(n):
          opt.zero_grad()
          loss = loss_fn(self(X_train), y_train)
          loss.backward()
          opt.step()
          
          losses.append(loss.item())
          
          if (i+1) % acc_step == 0:
            val, train_pred = torch.max(self(X_train), dim=1)
            val, test_pred  = torch.max(self(X_test), dim=1)
            
            train_acc.append( (train_pred == y_train).sum() / len(y_train) )
            test_acc.append( (test_pred == y_test).sum() / len(y_test) )
            
      return (losses, train_acc, test_acc)
```
:::


## Performance - relu

::: {.small}
```{python}
loss, train_acc, test_acc = mnist_fnn2_model(
  64,64,10, nl_step=torch.nn.ReLU()
).fit(
  X_train, y_train, X_test, y_test, n=1000
)
train_acc[-5:]
test_acc[-5:]
```
:::

```{python}
#| echo: false
#| out.width: 75%
plt.figure(figsize=(12,6))
plt.subplot(121)
plt.plot(loss, label="loss")
plt.legend()

plt.subplot(122)
plt.plot(range(0,1000,10), train_acc, label="train accuracy")
plt.plot(range(0,1000,10), test_acc, label="test accuracy")
plt.legend()

plt.show()
```


## Performance - tanh

::: {.small}
```{python}
loss, train_acc, test_acc = mnist_fnn2_model(
  64,64,10, nl_step=torch.nn.Tanh()
).fit(
  X_train, y_train, X_test, y_test, n=1000
)
train_acc[-5:]
test_acc[-5:]
```
:::

```{python}
#| echo: false
#| out.width: 75%
plt.figure(figsize=(12,6))
plt.subplot(121)
plt.plot(loss, label="loss")
plt.legend()

plt.subplot(122)
plt.plot(range(0,1000,10), train_acc, label="train accuracy")
plt.plot(range(0,1000,10), test_acc, label="test accuracy")
plt.legend()

plt.show()
```

# Convolutional NN

## 2d convolutions

::: {.aside}
Animation [source](https://towardsdatascience.com/intuitively-understanding-convolutions-for-deep-learning-1f6f42faee1)
:::

:::: {.columns .small}
::: {.column width='50%'}
```{r}
#| echo: false
knitr::include_graphics("imgs/tds_2dconv.gif")
```
:::

::: {.column width='50%' .fragment}
```{r}
#| echo: false
knitr::include_graphics("imgs/tds_2dconv2.gif")
```
:::
::::

## `nn.Conv2d()`

::: {.small}
```{python}
cv = torch.nn.Conv2d(
  in_channels=1, out_channels=4, 
  kernel_size=3, 
  stride=1, padding=1
)
```
:::

. . .

:::: {.columns .small}
::: {.column width='50%'}
```{python}
list(cv.parameters())[0] # kernel weights
```
:::

::: {.column width='50%'}
```{python}
list(cv.parameters())[1] # biases
```
:::
::::


## Applying `Conv2d()`

::: {.small}
```{python}
#| error: true
X_train[[0]]
X_train[[0]].shape
```
:::

. . .

::: {.small}
```{python}
#| error: true
cv(X_train[[0]])
```
:::

. . .

::: {.small}
```{python}
#| error: true
cv(X_train[[0]].view(1,8,8))
```
:::

## Pooling

::: {.small}
```{python}
x = torch.tensor(
  [[[0,0,0,0],
    [0,1,2,0],
    [0,3,4,0],
    [0,0,0,0]]],
  dtype=torch.float
)
x.shape
```
:::

. . .

:::: {.columns .small}
::: {.column width='50%'}
```{python}
p = torch.nn.MaxPool2d(kernel_size=2, stride=1)
p(x)

p = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
p(x)
```
:::

::: {.column width='50%' .fragment}
```{python}
p = torch.nn.AvgPool2d(kernel_size=2)
p(x)

p = torch.nn.AvgPool2d(kernel_size=2, padding=1)
p(x)
```
:::
::::

## Convolutional model

::: {.small}
```{python}
#| code-line-numbers: "|4-10,13-17"
class mnist_conv_model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn  = torch.nn.Conv2d(
          in_channels=1, out_channels=8,
          kernel_size=3, stride=1, padding=1
        )
        self.relu = torch.nn.ReLU()
        self.pool = torch.nn.MaxPool2d(kernel_size=2)
        self.lin  = torch.nn.Linear(8 * 4 * 4, 10)
        
    def forward(self, X):
        out = self.cnn(X.view(-1, 1, 8, 8))
        out = self.relu(out)
        out = self.pool(out)
        out = self.lin(out.view(-1, 8 * 4 * 4))
        return out
    
    def fit(self, X_train, y_train, X_test, y_test, lr=0.001, n=1000, acc_step=10):
      loss_fn = torch.nn.CrossEntropyLoss()
      opt = torch.optim.SGD(self.parameters(), lr=lr, momentum=0.9) 
      losses, train_acc, test_acc = [], [], []
      
      for i in range(n):
          opt.zero_grad()
          loss = loss_fn(self(X_train), y_train)
          loss.backward()
          opt.step()
          
          losses.append(loss.item())
          
          if (i+1) % acc_step == 0:
            val, train_pred = torch.max(self(X_train), dim=1)
            val, test_pred  = torch.max(self(X_test), dim=1)
            
            train_acc.append( (train_pred == y_train).sum() / len(y_train) )
            test_acc.append( (test_pred == y_test).sum() / len(y_test) )
            
      return (losses, train_acc, test_acc)
```
:::




## Performance

::: {.small}
```{python}
loss, train_acc, test_acc = mnist_conv_model().fit(
  X_train, y_train, X_test, y_test, n=1000
)
train_acc[-5:]
test_acc[-5:]
```
:::

```{python}
#| echo: false
#| out.width: 75%
plt.figure(figsize=(12,6))
plt.subplot(121)
plt.plot(loss, label="loss")
plt.legend()

plt.subplot(122)
plt.plot(range(0,1000,10), train_acc, label="train accuracy")
plt.plot(range(0,1000,10), test_acc, label="test accuracy")
plt.legend()

plt.show()
```



## Cleaning up models

::: {.small}
```{python}
#| code-line-numbers: "|4-14,17|5,12"
class mnist_conv_model2(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = torch.nn.Sequential(
          torch.nn.Unflatten(1, (1,8,8)),
          torch.nn.Conv2d(
            in_channels=1, out_channels=8,
            kernel_size=3, stride=1, padding=1
          ),
          torch.nn.ReLU(),
          torch.nn.MaxPool2d(kernel_size=2),
          torch.nn.Flatten(),
          torch.nn.Linear(8 * 4 * 4, 10)
        )
        
    def forward(self, X):
        return self.model(X)
    
    def fit(self, X_train, y_train, X_test, y_test, lr=0.001, n=1000, acc_step=10):
      opt = torch.optim.SGD(self.parameters(), lr=lr, momentum=0.9) 
      losses, train_acc, test_acc = [], [], []
      
      for i in range(n):
          opt.zero_grad()
          loss = torch.nn.CrossEntropyLoss()(self(X_train), y_train)
          loss.backward()
          opt.step()
          
          losses.append(loss.item())
          
          if (i+1) % acc_step == 0:
            val, train_pred = torch.max(self(X_train), dim=1)
            val, test_pred  = torch.max(self(X_test), dim=1)
            
            train_acc.append( (train_pred == y_train).sum() / len(y_train) )
            test_acc.append( (test_pred == y_test).sum() / len(y_test) )
            
      return (losses, train_acc, test_acc)
```
:::

# A bit more on non-linear<br/>activation layers

## Non-linear functions

```{python}
df = pd.read_csv("data/gp.csv")
X = torch.tensor(df["x"], dtype=torch.float32).reshape(-1,1)
y = torch.tensor(df["y"], dtype=torch.float32)
```

```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plt.show()
```

## Linear regression

::: {.small}
```{python}
class lin_reg(torch.nn.Module):
    def __init__(self, X):
        super().__init__()
        self.n = X.shape[0]
        self.p = X.shape[1]
        self.model = torch.nn.Sequential(
          torch.nn.Linear(self.p, self.p)
        )
    
    def forward(self, X):
        return self.model(X)
    
    def fit(self, X, y, n=1000):
      losses = []
      opt = torch.optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
      for i in range(n):
          loss = torch.nn.MSELoss()(self(X).squeeze(), y)
          loss.backward()
          opt.step()
          opt.zero_grad()
          losses.append(loss.item())
      
      return losses
```
:::

## Model results

::: {.small}
```{python}
m1 = lin_reg(X)
loss = m1.fit(X,y, n=2000)
```
:::

:::: {.columns .small}
::: {.column width='50%'}
Training loss:
```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
plt.plot(loss, label="loss")
plt.legend()
plt.show()
```
:::

::: {.column width='50%'}
Predictions
```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plt.plot(
  X.numpy().flatten(),
  m1(X).detach().numpy().flatten(),
  "-r",
  label="lin_reg"
)
plt.legend()
plt.show()
```
:::
::::



## Double linear regression

::: {.small}
```{python}
class dbl_lin_reg(torch.nn.Module):
    def __init__(self, X, hidden_dim=10):
        super().__init__()
        self.n = X.shape[0]
        self.p = X.shape[1]
        self.model = torch.nn.Sequential(
          torch.nn.Linear(self.p, hidden_dim),
          torch.nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, X):
        return self.model(X)
    
    def fit(self, X, y, n=1000):
      losses = []
      opt = torch.optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
      for i in range(n):
          loss = torch.nn.MSELoss()(self(X).squeeze(), y)
          loss.backward()
          opt.step()
          opt.zero_grad()
          losses.append(loss.item())
      
      return losses
```
:::

## Model results

::: {.small}
```{python}
m2 = dbl_lin_reg(X, hidden_dim=10)
loss = m2.fit(X,y, n=2000)
```
:::

:::: {.columns .small}
::: {.column width='50%'}
Training loss:
```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
plt.plot(loss, label="loss")
plt.legend()
plt.show()
```
:::

::: {.column width='50%'}
Predictions
```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plt.plot(
  X.numpy().flatten(),
  m2(X).detach().numpy().flatten(),
  "-g"
)
plt.legend()
plt.show()
```
:::
::::


## Non-linear regression w/ ReLU

::: {.small}
```{python}
class lin_reg_relu(torch.nn.Module):
    def __init__(self, X, hidden_dim=100):
        super().__init__()
        self.n = X.shape[0]
        self.p = X.shape[1]
        self.model = torch.nn.Sequential(
          torch.nn.Linear(self.p, hidden_dim),
          torch.nn.ReLU(),
          torch.nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, X):
        return self.model(X)
    
    def fit(self, X, y, n=1000):
      losses = []
      opt = torch.optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
      for i in range(n):
          loss = torch.nn.MSELoss()(self(X).squeeze(), y)
          loss.backward()
          opt.step()
          opt.zero_grad()
          losses.append(loss.item())
      
      return losses
```
:::

## Model results

::: {.small}
```{python}
#| echo: false
m3 = lin_reg_relu(X, hidden_dim=100)
loss = m3.fit(X,y, n=5000)
```

```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plt.plot(
  X.numpy().flatten(),
  m3(X).detach().numpy().flatten(),
  "-c"
)
plt.legend()
plt.show()
```
:::

## Hidden dimensions

```{python}
#| echo: false

m3_10   = lin_reg_relu(X, hidden_dim=10); z=m3_10.fit(X,y, n=2000)
m3_100  = lin_reg_relu(X, hidden_dim=100); z=m3_100.fit(X,y, n=2000)
m3_1000 = lin_reg_relu(X, hidden_dim=1000); z=m3_1000.fit(X,y, n=2000)

def plot_m(m, fmt, label):
  plt.plot(
    X.numpy().flatten(),
    m(X).detach().numpy().flatten(),
    fmt,
    label=label
  )


plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plot_m(m3_10, "-m", "hidden_dim=10")
plot_m(m3_100, "-g", "hidden_dim=100")
plot_m(m3_1000, "-r", "hidden_dim=1000")
plt.legend()
plt.show()
```



## Non-linear regression w/ Tanh

::: {.small}
```{python}
class lin_reg_tanh(torch.nn.Module):
    def __init__(self, X, hidden_dim=10):
        super().__init__()
        self.n = X.shape[0]
        self.p = X.shape[1]
        self.model = torch.nn.Sequential(
          torch.nn.Linear(self.p, hidden_dim),
          torch.nn.Tanh(),
          torch.nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, X):
        return self.model(X)
    
    def fit(self, X, y, n=1000):
      losses = []
      opt = torch.optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
      for i in range(n):
          loss = torch.nn.MSELoss()(self(X).squeeze(), y)
          loss.backward()
          opt.step()
          opt.zero_grad()
          losses.append(loss.item())
      
      return losses
```
:::


## Tanh & hidden dimension


```{python}
#| echo: false

m4_10   = lin_reg_tanh(X, hidden_dim=10); z=m4_10.fit(X,y, n=5000)
m4_100  = lin_reg_tanh(X, hidden_dim=100); z=m4_100.fit(X,y, n=5000)
m4_1000 = lin_reg_tanh(X, hidden_dim=1000); z=m4_1000.fit(X,y, n=5000)

plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plot_m(m4_10, "-m", "hidden_dim=10")
plot_m(m4_100, "-g", "hidden_dim=100")
plot_m(m4_1000, "-r", "hidden_dim=1000")
plt.legend()
plt.show()
```


## Three layers

::: {.small}
```{python}
class three_layers(torch.nn.Module):
    def __init__(self, X, hidden_dim=100):
        super().__init__()
        self.n = X.shape[0]
        self.p = X.shape[1]
        self.model = torch.nn.Sequential(
          torch.nn.Linear(self.p, hidden_dim),
          torch.nn.ReLU(),
          torch.nn.Linear(hidden_dim, hidden_dim),
          torch.nn.ReLU(),
          torch.nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, X):
        return self.model(X)
    
    def fit(self, X, y, n=1000):
      losses = []
      opt = torch.optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
      for i in range(n):
          loss = torch.nn.MSELoss()(self(X).squeeze(), y)
          loss.backward()
          opt.step()
          opt.zero_grad()
          losses.append(loss.item())
      
      return losses
```
:::

## Model results

```{python}
#| echo: false
m5 = three_layers(X, hidden_dim=100)
loss = m5.fit(X,y, n=10000)
```

::: {.small}
```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plt.plot(
  X.numpy().flatten(),
  m5(X).detach().numpy().flatten(),
  "-c"
)
plt.legend()
plt.show()
```
:::


## Five layers

::: {.small}
```{python}
class five_layers(torch.nn.Module):
    def __init__(self, X, hidden_dim=100):
        super().__init__()
        self.n = X.shape[0]
        self.p = X.shape[1]
        self.model = torch.nn.Sequential(
          torch.nn.Linear(self.p, hidden_dim),
          torch.nn.ReLU(),
          torch.nn.Linear(hidden_dim, hidden_dim),
          torch.nn.ReLU(),
          torch.nn.Linear(hidden_dim, hidden_dim),
          torch.nn.ReLU(),
          torch.nn.Linear(hidden_dim, hidden_dim),
          torch.nn.ReLU(),
          torch.nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, X):
        return self.model(X)
    
    def fit(self, X, y, n=1000):
      losses = []
      opt = torch.optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
      for i in range(n):
          loss = torch.nn.MSELoss()(self(X).squeeze(), y)
          loss.backward()
          opt.step()
          opt.zero_grad()
          losses.append(loss.item())
      
      return losses
```
:::

## Model results

::: {.small}
```{python}
#| echo: false
m6 = five_layers(X, hidden_dim=100)
loss = m6.fit(X,y, n=10000)
```
:::

```{python}
#| echo: false
plt.figure(figsize=(10,5), layout="constrained")
sns.scatterplot(data=df, x="x", y="y")
plt.plot(
  X.numpy().flatten(),
  m6(X).detach().numpy().flatten(),
  "-c"
)
plt.legend()
plt.show()
```
:::
