#导入需要的包

import pandas as pd

import numpy as np

import statsmodels.api as sm

import matplotlib.pyplot as plt

%matplotlib inline

#读入数据

data = pd.read_csv("Wage.csv")

data.head()

data_x = data['age']

data_y = data['wage']

#将数据划分为训练集和验证集

from sklearn.model_selection import train_test_split

train_x, valid_x, train_y, valid_y = train_test_split(data_x, data_y, test_size=0.33, random_state = 1)

#对年龄和工资的关系进行可视化

import matplotlib.pyplot as plt

plt.scatter(train_x, train_y, facecolor='None', edgecolor='k', alpha=0.3)plt.show()

rom sklearn.linear_model import LinearRegression

#拟合线性回归模型
x = train_x.reshape(-1,1)
model = LinearRegression()
model.fit(x,train_y)
print(model.coef_)
print(model.intercept_)
-> array([0.72190831])
-> 80.65287740759283
#在验证集上进行预测
valid_x = valid_x.reshape(-1,1)
pred = model.predict(valid_x)

#可视化
#我们将使用valid_x的最小值和最大值之间的70个点进行绘制
xp = np.linspace(valid_x.min(),valid_x.max(),70)
xp = xp.reshape(-1,1)
pred_plot = model.predict(xp)
plt.scatter(valid_x, valid_y, facecolor='None', edgecolor='k', alpha=0.3)
plt.plot(xp, pred_plot)
plt.show()

现在对预测出的结果算一下RMSE:
from sklearn.metrics import mean_squared_error
from math import sqrt
rms = sqrt(mean_squared_error(valid_y, pred))
print(rms)

#为回归函数生成权重，设degree=2

weights = np.polyfit(train_x, train_y, 2)

print(weights)
-> array([ -0.05194765,   5.22868974, -10.03406116])
#根据给定的权重生成模型
model = np.poly1d(weights)
#在验证集上进行预测
pred = model(valid_x)
#我们只画出其中的70个点
xp = np.linspace(valid_x.min(),valid_x.max(),70)
pred_plot = model(xp)
plt.scatter(valid_x, valid_y, facecolor='None', edgecolor='k', alpha=0.3)
plt.plot(xp, pred_plot)
plt.show()