首页 > 技术文章 > 用Python进行线性回归

popodynasty 2021-03-10 18:03 原文

from sklearn import linear_model
from matplotlib import pyplot as plt
import numpy as np

raw_xtrain = [8, 3, 9, 7, 16, 5, 13, 10, 4, 6]
ytrain = [30, 21, 33, 27, 42, 24, 36, 33, 22, 25]

raw_xtest = [4.5, 6, 14]
ytest = [24, 25, 38]

xtrain = []
xtest = []

for i in raw_xtrain:
    xtrain.append([i])
    
for i in raw_xtest:
    xtest.append([i])
    
clf = linear_model.LinearRegression()
clf.fit(xtrain, ytrain)
y_pred = clf.predict(xtest)

k = clf.coef_
b = clf.intercept_

x0 = np.arange(0, 20, 0.2)
y0 = k* x0 + b
# y = kx + b

mt = k * np.array(raw_xtrain) + b - np.array(ytrain)
mt = np.multiply(mt, mt)

print('train MSE={:.2f}'.format(1/10 * sum(mt)))

mt =k * np.array(raw_xtest) + b - np.array(ytest)
mt = np.multiply(mt, mt)

print('test MSE={:.2f}'.format(1/3 * sum(mt)))

print('figure:')
plt.title('LinearRegression')
plt.xlabel('X')
plt.ylabel('Y')
plt.grid()
plt.plot(raw_xtrain, ytrain, 'bo', label='train data')
plt.plot(raw_xtest, ytest, 'ro', label='test data')
plt.plot(raw_xtest, y_pred, 'go', label='pred')
plt.plot(x0, y0, 'black',label='pred line')
plt.legend()
plt.show()


# train MSE=0.92
# test MSE=0.67

推荐阅读