TensorFlow vs PyTorch — Linear Regression

Implementation of Linear Regression in both TensorFlow and PyTorch frameworks and comparison of their results

If your are new to Deep Learning and Neural Networks, then you must have come across the terms “TensorFlow” and “PyTorch”. These are two popular Deep Learning frameworks that are used in the field of Data Science.

In this exercise, I will be showing you the implementation of the most simple Neural Network (Linear Regression) using both the frameworks and compare their results.

Origin-

PyTorch is an open source machine learning library based on the Torch library. PyTorch was primarily developed by Facebook’s AI Research lab (FAIR). It is a free and open-source software.

On the other hand, TensorFlow was developed by the Google Brain team for internal Google research purpose. It is widely used for machine learning applications such as neural networks. It is also a free and open-source software.

TensorFlow vs PyTorch (Image from Source)

The most productive way of comparing two frameworks is to solve the same problem using both of them and analyzing their results. In this exercise, we shall perform Linear Regression using both TensorFlow and PyTorch frameworks and compare their results.

Problem-

We will be using a very simple example in this exercise. Here, we are given an array of numbers, x= [-1.0, 0.0, 1.0, 2.0, 3.0, 4.0] and y= [-3.0, -1.0, 1.0, 3.0, 5.0, 7.0]. The formula used in this is y = 2*x -1 , which is a Linear Regression.

x= [-1.0, 0.0, 1.0, 2.0, 3.0, 4.0]
y= [-3.0, -1.0, 1.0, 3.0, 5.0, 7.0]
y = 2*x -1

On training a Linear Regression model using both the frameworks, we shall input a new value for x=10 and see what the model predicts for y.

TensorFlow framework-

First, we shall go through the TensorFlow framework.

TensorFlow — Linear Regression

This is the code to perform Linear Regression with TensorFlow using keras library. Let us go through each block of code in the above program.

In the first step, we will import the libraries.

import tensorflow as tf
import numpy as np
from tensorflow import keras

In the next step, we design our model using the Sequential, which is a linear stack of layers. In this model, we use only one layer(neuron).

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

In the second step, we define our optimizer and loss functions to train our Neural Network model. In this, we use the Stochastic Gradient Descent (SDG) optimizer and the Mean Squared Error (MSE) as our loss function.

model.compile(optimizer='sgd', loss='mean_squared_error')

After this step, we initialize our numbers with two variables “xs” and “ys”.

xs = np.array([-1.0,  0.0, 1.0, 2.0, 3.0, 4.0], dtype=float)
ys = np.array([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0], dtype=float)

In the last step, we fit our model to the variables “xs” and “ys”. We train the model with 500 epochs.

model.fit(xs, ys, epochs=500)
>>Epoch 500/500
1/1 [==============================] - 0s 1ms/step - loss: 5.1584e-05

In this final step, we predict the value for a new xs = 10. According to the formula, y=2*x-1, for xs=10, we get a value of 19. Let us now see what value our model created using the TensorFlow framework predicts.

print(model.predict([10.0]))
>>[[18.979048]]

We see that, we get a value close to 19 using the Linear Regression model designed using the TensorFlow Framework.

PyTorch Framework-

Let us now see the Linear Regression Model designed using PyTorch framework.

PyTorch — Linear Regression

The PyTorch application of Linear Regression is indeed large and complex in comparison to the TensorFlow model. Let us analyze each step of the model.

In the first step, We import the required libraries for designing the Linear Regression model.

import torch
from torch.autograd import Variable

In the next step, we initialize the same values of “xs” and “ys” as defined in the TensorFlow model. For the PyTorch application, we convert the list to a Tensor using the appropriate functions.

xs = [[-1.0],  [0.0], [1.0], [2.0], [3.0], [4.0]]
ys = [[-3.0], [-1.0], [1.0], [3.0], [5.0], [7.0]]
xs = Variable(torch.Tensor(xs))
ys = Variable(torch.Tensor(ys))

After this, we define a class ‘LinearRegressionModel’ that we will use to define our model. Since this is a Simple Linear Regression with 1 input and 1 output, we use a Linear model with both the input and output dimension equal to 1. Finally, we create a “model” using the above-defined class.

class LinearRegressionModel(torch.nn.Module): 
  
    def __init__(self): 
        super(LinearRegressionModel, self).__init__() 
        self.linear = torch.nn.Linear(1, 1)  # One in and one out 
  
    def forward(self, x): 
        y_pred = self.linear(x) 
        return y_pred 

model = LinearRegressionModel()

Next, we choose our optimizer and loss criterion. We chose the same functions as used for the TensorFlow application, i.e, the SDG function for optimizer and MSE for the loss function. Additionally, we arbitrarily fix a learning rate of 0.01.

criterion = torch.nn.MSELoss(size_average = False) 
optimizer = torch.optim.SGD(model.parameters(), lr = 0.01)

We now arrive at our training step. In this stage, we perform three tasks for 500 iterations as we set the epoch value as 500.

1. Do a Forward pass by passing data and predicting the value of ys for each xs.
2. Computing the loss using MSE loss function.
3. Reset all gradients to 0, perform back-propagation finally updating the weights.

for epoch in range(500): 
  
    # Forward pass: Compute predicted y by passing  
    # x to the model 
    pred_y = model(xs) 
  
    # Compute and print loss 
    loss = criterion(pred_y, ys) 
  
    # Zero gradients, perform a backward pass,  
    # and update the weights. 
    optimizer.zero_grad() 
    loss.backward() 
    optimizer.step() 
    print('epoch {}, loss {}'.format(epoch, loss.item()))
>>epoch 499, loss 5.151434834260726e-13

Finally, we predict the value for a new xs=10 using this model. As described earlier, we must get a value close to 19. Let us now see what the PyTorch model outputs.

new_var = Variable(torch.Tensor([[10.0]])) 
pred_y = model(new_var) 
print(model(new_var).item())
>>18.999998092651367

We see the the Linear Regression model built using the PyTorch framework also gives us a value that is close to 19.

Comparision-

On visualizing the results obtained from both TensorFlow and PyTorch models, we see that the TensorFlow model gave us a result of 18.979048 and the PyTorch model gave us a result of 18.999998092651367.

TensorFlow Model — 18.979048

PyTorch Model — 18.999998092651367

It is clearly seen that a very simple Neural Network built using the PyTorch framework has more accuracy than the model built using the TensorFlow framework. However, the PyTorch model is more complex in nature and difficult to understand for a beginner.

If you are completely new to Deep Learning and Neural Networks, I would suggest you all to start with the TensorFlow framework and then move to the PyTorch framework once you gain experience in the former.

I have attached a detailed version of the codes (.ipynb) for both the frameworks in my Github profile for your reference.

mk-gurucharan/TensorFlow-vs-PyTorch-Linear-Regression-

I hope that I was able to explain and demonstrate the implementation of a simple neural network (Linear Regression Model) using the two most popular frameworks that are used today in Deep Learning. Till then, Happy Machine Learning!