How to Estimate Salary with Linear Regression
By: Bee Guan Teo
Tags: Getting Started, Programming, Data Science, Statistics, Python

Linear Regression

1. Loading data

Firstly, we will use the Python Pandas library to read our CSV data.

Line 1–6: Import all the required libraries.

Line 8: Use the Pandas read_csv function to read the CSV file. This function will return the data in a dataframe format.

Image Prepared by the Author

Line 9–10: Extract the column of YearsExperience and Salary and assign them to the variables X and y, respectively.

2. Splitting data into a training set and a test set

Prior to building a linear model, we need to prepare a training set and a test set (part of standard procedure in a machine learning workflow). The training set will be used to train the model whereas the test set will be used to assess the performance of the trained model in predicting the result from unseen data.

We can use the Python scikit-learn train_test_split function to randomly split our data into a training and test set.

Line 1: Set apart 30% of the entire dataset as the test set and assign the training set and test set into four variables, respectively.

3. Data Transformation

Python scikit-learn only accepts the training and test data in a 2-dimensional array format. We have to perform data transformation on our training set and test set.

Line 1–2: Use Numpy reshape function to transform the training set from 1-dimensional series to a 2-dimensional array.

Line 3–4: Use Numpy reshape function to transform the test set from 1-dimensional series to a 2-dimensional array.

4. Training Model

Now we are ready to train our linear model.

Line 1: Use the Scikit-Learn LinearRegression function to create a model object.

Line 2: Fit the training set to the model.

5. Predicting Salary using Linear Model

At this stage, we have trained a linear model and we first use it to predict the salary on our training set to see how well it fit on the data.

Line 1: Use the linear model to predict the salary based on the training set.

Line 3–10: Use the Matplotlib to create a plot to visualize the predicted results. The “true values” are plotted as the blue dots on the chart and the predicted values are plotted as a black color straight line.

Image Prepared by the Author

In general, the linear model fits well on the training data. This shows a linear relationship between the salary and years of experience.

Now, we need to check if the linear model can perform well on our test set (unknown data).

Line 1: Use the linear model to predict the salary based on the test set.

Line 3–10: Use the Matplotlib to create a plot to visualize the predicted results. The “true values” are plotted as the green dots on the chart and the predicted values are plotted as a black color straight line.

Image Prepared by the Author

The graph shows that our linear model can fit quite well on the test set. We can observe a linear pattern of how the amount of salary is increased by the years of experience.

6. Model Evaluation

The previous section uses a graphical approach to evaluate the performance of our linear model which can sometimes be quite subjective in our judgment. Here we will use some quantitative methods to obtain a more precise performance evaluation of our linear model.

We will use three types of quantitative metrics:

Mean Square Error — The average of the squares of the difference between the true values and the predicted values. The lower the difference the better the performance of the model. This is a common metric used for regression analysis.

— The average of the squares of the difference between the true values and the predicted values. The lower the difference the better the performance of the model. This is a common metric used for regression analysis. Explained Variance Score — A measurement to examine how well a model can handle the variation of values in the dataset. A score of 1.0 is the perfect score.

— A measurement to examine how well a model can handle the variation of values in the dataset. A score of 1.0 is the perfect score. R2 Score — A measurement to examine how well our model can predict values based on the test set (unknown samples). The perfect score is 1.0.

Line 1–3: Use the Scikit-learn Metrics functionalities to calculate the mean square error, explained variance, and R2 score for our linear model. We feed the functions with the true values (test set) and the predicted values.

Image Prepared by the Author

If we perform a root square calculation on our mean squared error, we will gain an average discrepancy of around $6146.92 which is quite a low error. Besides, the explain variance score and the R2 score hit 0.9 above. This shows our linear model is not overfitted and can work nicely to predict the salary based on new data.