34 lines
1 KiB
Python
34 lines
1 KiB
Python
# Data preprocessing
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.linear_model import LinearRegression
|
|
|
|
# Importing the dataset
|
|
dataset = pd.read_csv('Salary_Data.csv')
|
|
X = dataset.iloc[:, :-1].values
|
|
y = dataset.iloc[:, 1].values
|
|
|
|
X_train, X_test, y_train, y_test = train_test_split(
|
|
X, y, test_size=0.2, random_state=0)
|
|
|
|
# Fitting Simple Linear Regreesion to the Training set
|
|
regressor = LinearRegression()
|
|
regressor.fit(X_train, y_train)
|
|
|
|
# Prediction the Test set result
|
|
# y_pred contains the predicted salary from the test sample, y_test is the actual salary.
|
|
y_pred = regressor.predict(X_test)
|
|
|
|
# Visualize the data
|
|
# data use to train the regression
|
|
plt.scatter(X_train, y_train, color='red')
|
|
# actual data we compared with our trained regression
|
|
plt.scatter(X_test, y_test, color='green')
|
|
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()
|