A Smooth Introduction to Linear Regression and its Implementation in PyTorch (Part-II)

So in Part-I I gave a simple introduction on what linear regression is and how we can find the equation of the best fit line for our data. In this post, I will show you how to implement the task we worked on in Part-I in PyTorch. So, let’s get started!

The first step is to import the libraries we will be using:

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

Now, we need to set the hyperparameters, as follows:

input_size = 1
output_size = 1
epochs = 50
learning_rate = 0.0001

The input size is set to 1 since our inputs to the model are scalars h (hour in the day). The output is also set to 1 since we will get only one value returned for r (number of pages being read). So, basically, we will leave our program to find the best values for B_0 and B_1 that we calculated in the previous part of this tutorial.

Let’s now enter the training data, which will represent the values of h and r that we already have.

h_train = np.array([[0],[1],[2],[3],[4],[5]],dtype=np.float32)
r_train = np.array([[10],[12],[7],[13],[5],[9]],dtype=np.float32)

To perform linear regression we have to define three things: model (linear regression), loss function, and the optimizer, after which we can take our data to training. The figure below depicts this process.

Let’s take that step-by-step in PyTorch. So, first we define our linear regression model:

model = nn.Linear(input_size,output_size)

Then, define the loss function (mean squared error):

criterion = nn.MSELoss()

And the optimizer (stochastic gradient descent):

optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate)

We can now train our model with the number of epochs specified (i.e. 5).

# train the model
for epoch in range(epochs):
	inputs = torch.from_numpy(h_train)
	targets = torch.from_numpy(r_train)
	# forward propagation
	predictions = model(inputs)
	loss = loss_function(predictions,targets)
	# backward propagation
	loss.backward()
	# optimization
	optimizer.step()

Oh, to get an idea on what’s going on after each epoch, we can add the following statement in our training for-loop:

print 'Epoch: ' + str(epoch) + '.....' + 'Loss: ' + str(loss.item())

That’s it! You have just written a PyTorch program that will find the best fit line (i.e. linear regression) for our data which describes how many pages the person read at each hour of the day.

Let’s go ahead and plot our best fit line against the original data we provided the model with, as follows:

prediction = model(torch.from_numpy(h_train)).detach().numpy()
plt.plot(h_train, r_train, 'v--g', label='original data')
plt.plot(h_train, prediction, 'r', label='best fit line')
plt.legend()
plt.show()

And,….. this is how our best fit line looks like!

If you would like the full code, you can kindly find it here.