sklearn怎么分别计算决策树、线性回归的训练集MSE、测试集MSE并得出训练过程和测试结果,以三千条数据为例
时间: 2024-05-03 15:18:26 浏览: 83
针对决策树和线性回归的训练集MSE、测试集MSE,我们可以采用交叉验证的方法来计算。在sklearn中,交叉验证可以使用KFold或者train_test_split方法实现。
以下是使用train_test_split方法计算决策树和线性回归的训练集MSE、测试集MSE的示例代码:
```python
from sklearn.tree import DecisionTreeRegressor
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
# 加载数据,假设数据已经存储在X和y数组中
X = ...
y = ...
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)
# 训练决策树模型
dt = DecisionTreeRegressor()
dt.fit(X_train, y_train)
# 计算训练集MSE
y_train_pred = dt.predict(X_train)
train_mse_dt = mean_squared_error(y_train, y_train_pred)
# 计算测试集MSE
y_test_pred = dt.predict(X_test)
test_mse_dt = mean_squared_error(y_test, y_test_pred)
# 输出决策树模型的训练集MSE和测试集MSE
print("决策树模型的训练集MSE: {:.3f}".format(train_mse_dt))
print("决策树模型的测试集MSE: {:.3f}".format(test_mse_dt))
# 训练线性回归模型
lr = LinearRegression()
lr.fit(X_train, y_train)
# 计算训练集MSE
y_train_pred = lr.predict(X_train)
train_mse_lr = mean_squared_error(y_train, y_train_pred)
# 计算测试集MSE
y_test_pred = lr.predict(X_test)
test_mse_lr = mean_squared_error(y_test, y_test_pred)
# 输出线性回归模型的训练集MSE和测试集MSE
print("线性回归模型的训练集MSE: {:.3f}".format(train_mse_lr))
print("线性回归模型的测试集MSE: {:.3f}".format(test_mse_lr))
```
以上代码中,我们使用train_test_split方法将原始数据划分为训练集和测试集(测试集占总数据量的30%),然后分别训练决策树模型和线性回归模型,并计算它们的训练集MSE和测试集MSE。
需要注意的是,以上代码只是一个示例,具体实现方式可能因数据集的不同而有所差异。另外,如果需要可视化训练过程和测试结果,可以使用matplotlib等工具进行绘图。
阅读全文