Gradient Descent using Pytorch

gradient_descent.py

import torch
import numpy as np

# f = w * x
# f = 2 * x

X = np.array([1,2,3,4,5],dtype=np.float32)
Y = np.array([5,10,15,20,25],dtype=np.float32)

w = 0.0

# model prediction 
def forward(x):
    return w * x

# loss 
def loss(y,y_predicted):
    return ((y_predicted - y)**2).mean()

# gradient

# MSE = 1/N * (w*x -y)**2
# dj/dw = 1/N 2x (wx -y)

def gradient(x,y,y_predicted):
    return np.dot(2*x,y_predicted-y).mean()

print(f"Prediction before training: f(5) = {forward(5):.3f}")

# Training
learning_rate =0.01
n_iter = 10

for epoch in range(n_iter):
    
    # prediction = forward pass
    y_pred = forward(X)
    
    # loss
    l = loss(Y, y_pred)
    
    # gradient
    dw = gradient(X,Y,y_pred)
    
    # update weights
    w -= learning_rate * dw
    
    if epoch % 1 == 0:
        print(f"epoch : {epoch+1}, w = {w:.3f}, loss = {l:.8f}")
        
print(f"Prediction after training: f(5) = {forward(5):.3f}")