simple linear regression - salary prediction

DS/ML: Simple Linear Regression implementation (Python and R)

Implementing salary prediction using simple linear regression in Python and R.

Both Python and R code are available in my Github page. You can download the sample csv file (Salary_Data.csv) from the same place.

Python Code:


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

# Importing the dataset
dataset = pd.read_csv('Salary_Data.csv')
X = dataset.iloc[:, :-1].values
y = dataset.iloc[:, 1].values

# Splitting the dataset into the Training set and Test set
from sklearn.cross_validation import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 1/3, random_state =0)

# Fitting Simple Linear Regression to the Training set
from sklearn.linear_model import LinearRegression
regressor = LinearRegression()
regressor.fit(X_train, y_train)

# Predicting the Test set results
y_pred = regressor.predict(X_test)

# Visualising the Training set results
plt.scatter(X_train, y_train, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title('Salary vs Experience (Training set)')
plt.xlabel('Years of Experience')
plt.ylabel('Salary')
plt.show()

# Visualising the Test set results
plt.scatter(X_test, y_test, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title('Salary vs Experience (Test set)')
plt.xlabel('Years of Experience')
plt.ylabel('Salary')
plt.show()

R code:


# Importing the dataset
dataset = read.csv('Salary_Data.csv')

# Splitting the dataset into the Training set and Test set
# install.packages('caTools')
library(caTools)
set.seed(123)
split = sample.split(dataset$Salary, SplitRatio = 2/3)
training_set = subset(dataset, split == TRUE)
test_set = subset(dataset, split == FALSE)

# Fitting Simple Linear Regression to the Training set
regressor = lm(formula = Salary ~ YearsExperience,
 data = training_set)

# Predicting the Test set results
y_pred = predict(regressor, newdata = test_set)

# Visualising the Training set results
library(ggplot2)
ggplot() +
 geom_point(aes(x = training_set$YearsExperience, y = training_set$Salary),
 colour = 'red') +
 geom_line(aes(x = training_set$YearsExperience, y = predict(regressor, newdata = training_set)),
 colour = 'blue') +
 ggtitle('Salary vs Experience (Training set)') +
 xlab('Years of experience') +
 ylab('Salary')

# Visualising the Test set results
library(ggplot2)
ggplot() +
 geom_point(aes(x = test_set$YearsExperience, y = test_set$Salary),
 colour = 'red') +
 geom_line(aes(x = training_set$YearsExperience, y = predict(regressor, newdata = training_set)),
 colour = 'blue') +
 ggtitle('Salary vs Experience (Test set)') +
 xlab('Years of experience') +
 ylab('Salary')

Output:

simple linear regression - salary prediction

Credits: Udemy

Leave a Reply

Your email address will not be published. Required fields are marked *