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

import torch

N, D_in, H, D_out = 64,1000,100,10

# +
x = torch.randn(N,D_in)
y = torch.randn(N,D_out)

w1 = torch.randn(D_in,H,requires_grad=True)
w2 = torch.randn(H,D_out,requires_grad=True)
# -

learning_rate = 1e-6
for t in range(500):
    # mm is similar with matmul, but not support broadcast
    y_pred = x.mm(w1).clamp(min=0).mm(w2)
    loss = (y_pred - y).pow(2).sum()
    
    if t % 100 == 0:
        print(loss)
        
    loss.backward()
    
    with torch.no_grad():
        w1 -= learning_rate*w1.grad
        w2 -= learning_rate*w2.grad
        w1.grad.zero_()
        w2.grad.zero_()

# ## PyTorch nn

# +
model = torch.nn.Sequential(
            torch.nn.Linear(D_in,H),
            torch.nn.ReLU(),
            torch.nn.Linear(H,D_out))

learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate)

for t in range(500):
    y_pred = model(x)
    loss = torch.nn.functional.mse_loss(y_pred,y)
    
    if t % 100 == 0:
        print(loss)
    loss.backward()
    
    optimizer.step()
    optimizer.zero_grad()
# -

# ## PyTorch DataLoaders

import torch
from torch.utils.data import DataLoader,TensorDataset
import numpy as np

# +
N, D_in, H, D_out = 64,1000,100,10
# x = torch.randn(N,D_in)
# y = torch.randn(N,D_out)

x = np.random.randn(N,D_in)
y = np.random.randn(N,D_out)
# -

x = torch.Tensor(x)
y = torch.Tensor(y)

loader = DataLoader(TensorDataset(x,y),batch_size=8)


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)   # clamp(min=0) is the same as ReLU function
        y_pred = self.linear2(h_relu)
        return y_pred


model = TwoLayerNet(D_in,H,D_out)
optimizer = torch.optim.SGD(model.parameters(),lr=1e-2)

for epoch in range(20):
    epoch_loss = 0
    for x_batch,y_batch in loader:
        y_pred = model(x_batch)
        loss = torch.nn.functional.mse_loss(y_pred,y_batch)
        
        epoch_loss += loss
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    print('epoch {}, loss:{}'.format(epoch,epoch_loss))





StarcoderdataJupyterScriptsDedupFiltered