69 lines
1.9 KiB
Python
69 lines
1.9 KiB
Python
# prediction pytorch model -> forward function gets replaced by pytroch model
|
|
# gradient computation autograd
|
|
# loss computation pytorch loss
|
|
# parameter update pytorch optimizer
|
|
|
|
import torch
|
|
import torch.nn as nn # neural network module
|
|
|
|
|
|
# linear regression, no bias
|
|
# f = w*x
|
|
# f = 2*x
|
|
|
|
X = torch.tensor([[1], [2], [3], [4]], dtype=torch.float32) #reshape for pytorch model
|
|
Y = torch.tensor([[2], [4], [6], [8]], dtype=torch.float32)
|
|
X_test = torch.tensor([5], dtype=torch.float32)
|
|
|
|
n_samples, n_features = X.shape
|
|
print(n_samples, n_features)
|
|
|
|
input_size = n_features
|
|
output_size = n_features
|
|
|
|
# model = nn.Linear(input_size, output_size, bias=False)
|
|
|
|
#custom linear regression model (just a wrapper in this case, but you can add more layers)
|
|
class LinearRegression(nn.Module):
|
|
def __init__(self, input_dim, output_dim, bias=True):
|
|
super(LinearRegression, self).__init__()
|
|
#define layers
|
|
self.lin = nn.Linear(input_dim, output_dim, bias)
|
|
|
|
def forward(self, x):
|
|
return self.lin(x)
|
|
|
|
model = LinearRegression(input_size, output_size, bias=False)
|
|
|
|
print(f'Prediction before training: f(5) = {model(X_test).item():.3f}')
|
|
|
|
#Training
|
|
learning_rate = .01
|
|
n_iters = 100
|
|
|
|
loss = nn.MSELoss() # use pytorch built in MSE loss function
|
|
optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate) # use pytorch built in optimizer to optimize parameter 'w' with learning rate
|
|
|
|
|
|
for epoch in range(n_iters):
|
|
# prediction = forward pass
|
|
y_pred = model(X)
|
|
|
|
# loss
|
|
l = loss(Y, y_pred)
|
|
|
|
# gradients = backward pass
|
|
l.backward()
|
|
|
|
#update weights
|
|
optimizer.step()
|
|
|
|
# clear gradients
|
|
optimizer.zero_grad()
|
|
|
|
if epoch % 10 == 0: #every nth epoch
|
|
[w] = model.parameters()
|
|
print(f'epoch {epoch+1}: w = {w[0].item():.3f}, loss = {l:.8f}')
|
|
|
|
print(f'Prediction after training: f(5) = {model(X_test).item():.3f}')
|