Simple Linear Regression Using TensorFlow Vs PyTorch

In this article, we’ll see how to implement Simple Linear Regression using TensorFlow Vs PyTorch.

Linear Regression is a method that seeks to find a linear association between a dependent variable and an independent variable by reducing the gap in between.

Import libraries need for Simple Linear Regression model such as NumPy for array manipulation, matplotlib for plotting result and pandas for loading data (Download Dataset).

Simple Linear Regression using TensorFlow

Import libraries

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

Load Dataset

dataset = pd.read_csv('Salary_Data.csv')
print(dataset)
X_train = dataset.iloc[:, :-1].values
y_train = dataset.iloc[:, -1].values

Create Model in TensorFlow

model = tf.keras.Sequential([
  tf.keras.layers.Dense(units=1, input_shape=[1])
])

Assign Loss Function and Optimiser

model.compile(loss='mean_squared_error',
              optimizer=tf.keras.optimizers.SGD())

Run Model

history = model.fit(X_train, y_train, epochs=2000)

Model Prediction

print(model.predict([2.0]))

Full Code

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

#Load Dataset
dataset = pd.read_csv('Salary_Data.csv')
X_train = dataset.iloc[:, :-1].values
y_train = dataset.iloc[:, -1].values

#Model
model = tf.keras.Sequential([
  tf.keras.layers.Dense(units=1, input_shape=[1])
])

# Loss Function And Optimizer 
model.compile(loss='mean_squared_error',
              optimizer=tf.keras.optimizers.SGD())
# Run Model
history = model.fit(X_train, y_train, epochs=2000)

#Plot Result
plt.scatter(X_train, y_train, color = 'red')
plt.plot(X_train, model.predict(X_train), color = 'blue')
plt.title('Salary vs Experience (Training set TensorFlow)')
plt.xlabel('Years of Experience')
plt.ylabel('Salary')
plt.show()

Simple Linear Regression using Pytorch

Import libraries

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

Load Dataset And Convert to Tensor

dataset = pd.read_csv('Salary_Data.csv')
print(dataset)
X_train = Variable(torch.tensor(dataset.iloc[:, :-1].values.astype(np.float32)).cuda())
y_train = Variable(torch.tensor(dataset.iloc[:, -1].values.astype(np.float32)).cuda())

Reshape Tensor Data

y_train = y_train.view(y_train.shape[0], 1)

Create Model

modelT = nn.Linear(1, 1).cuda()

Assign Loss Function and Optimiser

learning_rate = 0.01
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(modelT.parameters(), lr=learning_rate)

Run Model

num_epochs = 2000
for epoch in range(num_epochs):
    # Forward pass and loss
    y_predicted = modelT(X_train)
    loss = criterion(y_predicted, y_train)
    
    # Backward pass and update
    loss.backward()
    optimizer.step()

    # zero grad before new step
    optimizer.zero_grad()

    if (epoch+1) % 10 == 0:
        print(f'epoch: {epoch+1}, loss = {loss.item():.4f}')

Model Prediction

new_var = Variable(torch.Tensor([[10.6]]).cuda()) 
print(modelT(new_var).item())

Full Code

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

dataset = pd.read_csv('Salary_Data.csv')
print(dataset)
X_train = Variable(torch.tensor(dataset.iloc[:, :-1].values.astype(np.float32)).cuda())
y_train = Variable(torch.tensor(dataset.iloc[:, -1].values.astype(np.float32)).cuda())

# Reshape Data
y_train = y_train.view(y_train.shape[0], 1)

# Model
modelT = nn.Linear(1, 1).cuda()

# Loss Function And Optimizer 
learning_rate = 0.01
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(modelT.parameters(), lr=learning_rate)

num_epochs = 2000
for epoch in range(num_epochs):
    # Forward pass and loss
    y_predicted = modelT(X_train)
    loss = criterion(y_predicted, y_train)
    
    # Backward pass and update
    loss.backward()
    optimizer.step()

    # zero grad before new step
    optimizer.zero_grad()

    if (epoch+1) % 10 == 0:
        print(f'epoch: {epoch+1}, loss = {loss.item():.4f}')

predicted = modelT(X_train).cpu().detach().numpy()
#Plot Data
plt.scatter(X_train.cpu(), y_train.cpu(), color = 'red')
plt.plot(X_train.cpu(), predicted, color = 'blue')
plt.title('Salary vs Experience (Training set Pytorch)')
plt.xlabel('Years of Experience')
plt.ylabel('Salary')
plt.show()