Linear Regression using Pytorch

LinearRegression.py

# 1) Design model (input, output size, forward pass)
# 2) Construct loss and optimizer 
# 3) Training loop
# - forward pass: compute prediction
# - backward pass: gradients
# - update weights

import torch
import torch.nn as nn
import numpy as np
from sklearn import datasets
import matplotlib.pyplot as plt

# prepare data
x_numpy, y_numpy = datasets.make_regression(n_samples=100,n_features=1,noise=20,random_state=1)
X = torch.from_numpy(x_numpy.astype(np.float32))
Y = torch.from_numpy(y_numpy.astype(np.float32))
Y = Y.view(Y.shape[0],1)
n_samples , n_features = X.shape
# model 
input_size = n_features
output_size = 1
model = nn.Linear(input_size, output_size)
# loss and optmizer
learning_rate = .0001
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate)
# training loop
n_iter = 10000
for epoch in range(n_iter):
    # forward
    y_pred = model(X)
    # loss
    loss = criterion(y_pred,Y)
    # backward
    loss.backward()
    # update
    optimizer.step()
    # zero grad
    optimizer.zero_grad()
    if (epoch +1) % 100 == 0:
        print(f"epoch: {epoch+1}, loss={loss.item():.4f}")
         
         
# plot
predicted = model(X).detach().numpy()
plt.plot(x_numpy,y_numpy,'ro')
plt.plot(x_numpy,predicted,'b')
plt.show()