1. 程式人生 > >一個簡單的線性迴歸例子

一個簡單的線性迴歸例子

from sklearn.model_selection 
import train_test_split
path = '8.Advertising.csv'
data = pd.read_csv(path)    # TV、Radio、Newspaper、Sales
x = data[['TV', 'Radio', 'Newspaper']]
y = data['Sales']
x_train, x_test, y_train, y_test = train_test_split(x, y, random_state=1)
linreg = LinearRegression()#獲取線性迴歸的模型
linreg.fit(x_train, y_train)#將資料填入模型中
y_hat = linreg.predict(np.array(x_test))#將x_test轉換成陣列來做預測
t = np.arange(len(x_test))
plt.plot(t, y_test, 'r-', linewidth=2, label='Test')
plt.plot(t, y_hat, 'g-', linewidth=2, label='Predict')
plt.show()