# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3.7
#     language: python
#     name: py37
# ---

# # 1. NeuralNetwork from scratch - numpy

import numpy as np

# N: batch size
# D_in: input dimension
# H: hidden dimension
# D_out: output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# +
# Create dummy input and output data
X = np.random.randn(N, D_in)
y = np.random.randn(N, D_out)

print(X.shape)
print(y.shape)


# -

# #### Model
# - Forward
#
# $$\begin{split}
# & h = X \cdot w_1 \\
# & \text{h_relu} = relu(h) \\
# & y\_ = relu(\text{h_relu} \cdot w_2)
# \end{split}$$
#
# - Loss function:
#
# $$loss = \sum\limits_{i=1}^n(y_i - \hat{y}_i)^2 = (y - y\_)^2.sum()$$
#
# - Backward
#
# $$\begin{split}
# & \text{grad_y_} = \frac{\delta\ loss}{\delta y\_} = -2(y - y\_) = 2(y\_ - y) \\
# & \text{grad_w}_2 = \frac{\delta\ loss}{\delta w_2} = \frac{\delta\ y\_}{\delta w_2}\ \frac{\delta\ loss}{\delta y\_} = \text{h_relu}^T \cdot \text{grad_y_} \\
# & \text{grad_h_relu} = \frac{\delta\ loss}{\delta \text{h_relu}} = \frac{\delta\ loss}{\delta y\_}\ \frac{\delta\ y\_}{\delta \text{h_relu}} = \text{grad_y_} \cdot w_2^T \\
# & \cdots \\
# & \text{grad_w}_1 = X^T \cdot \text{grad_h}
# \end{split}$$

def relu(x): return np.maximum(0, x)


# +
# %%time
# Randomly initialize params
w1 = np.random.randn(D_in, H)
w2 = np.random.randn(H, D_out)
learning_rate = 1e-6

for epoch in range(1, 501):
    # Forward pass: compute y_
    h = np.dot(X, w1)
    h_relu = relu(h)
    y_ = np.dot(h_relu, w2)

    # Compute and print loss
    loss = np.square(y_ - y).sum()
    if epoch % 100 == 0: print(epoch, loss)

    # Backprop to compute gradients of w1 and w2 with respect to loss
    grad_y_ = 2.0 * (y_ - y)
    grad_w2 = np.dot(h_relu.T, grad_y_)

    grad_h_relu = np.dot(grad_y_, w2.T)
    grad_h = grad_h_relu.copy()
    grad_h[h < 0] = 0
    grad_w1 = np.dot(X.T, grad_h)

    # Update weights
    w1 -= learning_rate * grad_w1
    w2 -= learning_rate * grad_w2
# -

# # 2. NeuralNetwork from scratch - torch tensor

# +
import torch

device = torch.device('cpu')
dtype=torch.float
# -

# N: batch size
# D_in: input dimension
# H: hidden dimension
# D_out: output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# +
# Create dummy input and output data
X = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

print(X.size())
print(y.size())

# +
# %%time
# Randomly initialize weights
w1 = torch.randn(D_in, H, device=device, dtype=dtype)
w2 = torch.randn(H, D_out, device=device, dtype=dtype)

learning_rate = 1e-6

for epoch in range(1, 501):
    # Forward
    h = X.mm(w1)
    h_relu = h.clamp(min=0)
    y_ = h_relu.mm(w2)
    
    # Compute and print loss
    loss = (y_ - y).pow(2).sum().item()
    if epoch % 100 == 0: print(epoch, loss)
    
    # Backprop to compute gradients of w1 and w2 with respect to loss
    grad_y_ = 2.0 * (y_ - y)
    grad_w2 = h_relu.t().mm(grad_y_)

    grad_h_relu = grad_y_.mm(w2.t())
    grad_h = grad_h_relu.clone()
    grad_h[h < 0] = 0
    grad_w1 = X.t().mm(grad_h)
    
    # Update weights
    w1 -= learning_rate * grad_w1
    w2 -= learning_rate * grad_w2
# -

# # 3. NeuralNetwork from scratch - torch Autograd

# +
import torch

device = torch.device('cuda')
dtype=torch.float
# -

# N: batch size
# D_in: input dimension
# H: hidden dimension
# D_out: output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# +
# Create dummy input and output data
X = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

print(X.size())
print(y.size())

# +
# %%time
# Randomly initialize weights
# requires_grad=True indicates that we want to compute gradients
w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)

learning_rate = 1e-6

for epoch in range(1, 501):
    # Forward
    h = X.mm(w1)
    h_relu = h.clamp(min=0)
    y_ = h_relu.mm(w2)
    
    # Compute and print loss
    loss = (y_ - y).pow(2).sum()
    if epoch % 100 == 0: print(epoch, loss.item())
    
    # Backprop to compute gradients of w1 and w2 with respect to loss
    loss.backward()

    # Update weights
    with torch.no_grad():
        w1 -= learning_rate * w1.grad
        w2 -= learning_rate * w2.grad
        
        # Manually zero the gradients after updating weights
        w1.grad.zero_()
        w2.grad.zero_()
# -

# # 4. NeuralNetwork with Pytorch nn

# +
import torch

device = torch.device('cpu')
dtype=torch.float
# -

# N: batch size
# D_in: input dimension
# H: hidden dimension
# D_out: output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# +
# Create dummy input and output data
X = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

print(X.size())
print(y.size())

# +
# %%time
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Linear(H, D_out),
)
loss_fn = torch.nn.MSELoss(reduction='sum')
learning_rate = 1e-6

for epoch in range(1, 501):
    # Forward
    y_ = model(X)
    
    # Compute and print loss
    loss = loss_fn(y_, y)
    if epoch % 100 == 0: print(epoch, loss.item())
    
    # Backprop to compute gradients of w1 and w2 with respect to loss
    model.zero_grad()
    loss.backward()

    # Update weights
    with torch.no_grad():
        for param in model.parameters():
            param -= learning_rate * param.grad
# -

# #### With PyTorch: optim

# +
# %%time
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Linear(H, D_out),
)
loss_fn = torch.nn.MSELoss(reduction='sum')
learning_rate = 1e-6
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(1, 501):
    # Forward
    y_ = model(X)
    
    # Compute and print loss
    loss = loss_fn(y_, y)
    if epoch % 100 == 0: print(epoch, loss.item())
    
    # Backprop to compute gradients of w1 and w2 with respect to loss
    optimizer.zero_grad()
    loss.backward()

    # Update weights
    optimizer.step()
# -

# # 5. NeuralNetwork with Custom nn

# +
import torch

device = torch.device('cpu')
dtype=torch.float
# -

# N: batch size
# D_in: input dimension
# H: hidden dimension
# D_out: output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# +
# Create dummy input and output data
X = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

print(X.size())
print(y.size())


# -

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        super(TwoLayerNet, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    def forward(self, X):
        h_relu = self.linear1(X).clamp(min=0)
        y_ = self.linear2(h_relu)
        return y_


# +
# %%time
model = TwoLayerNet(D_in, H, D_out)
loss_fn = torch.nn.MSELoss(reduction='sum')
learning_rate = 1e-6
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(1, 501):
    # Forward
    y_ = model(X)
    
    # Compute and print loss
    loss = loss_fn(y_, y)
    if epoch % 100 == 0: print(epoch, loss.item())
    
    # Backprop to compute gradients of w1 and w2 with respect to loss
    optimizer.zero_grad()
    loss.backward()

    # Update weights
    optimizer.step()

StarcoderdataJupyterScriptsDedupFiltered