《scikit-learn》决策树之回归树

今天我们来看看怎么使用回归树，其实大致流程还是跟决策树一样的，只不过我们今天要使用的是tree.DecisionTreeRegressor这个类。

1：加载数据
由于是要做回归，因此我们需要连续的数据，比如拿波士顿房价距离。

from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
from sklearn import tree
import graphviz

boston = load_boston()   #这里使用boston数据集
print(boston.data.shape)
print(boston.target.shape)
print(boston.feature_names)  # 所有特征属性的名称

data_train, data_test, target_train, target_test = train_test_split(boston.data, boston.target, test_size=0.3)
print(data_train.shape)
print(target_train.shape)
print(data_test.shape)
print(target_test.shape)

2：实例化回归树模型，进行训练

# 训练模型，用score对训练好的模型用测试数据进行评估。
regressor = tree.DecisionTreeRegressor(criterion='mse',
                                       random_state=0,
                                       min_samples_leaf=2,
                                       min_samples_split=10)
regressor = regressor.fit(data_train, target_train)  # 用训练数据和标签，进行训练。简单的就这么简单就训练出了模型。
score = regressor.score(data_test, target_test)  # R平方评估
print(score)

如果觉得不过瘾，想看看到底运行成功咋样，我们使用predict函数来进行预测。

# 单独挑前10个数据进行预测试一试，拎出来看看效果。
print(target_test[0:10])
predict = regressor.predict(data_test)  # 返回每个测试样本分类/回归的结果
print(predict[0:10])

从这里我们可以看到实际对数据的回归预测

3：模型回归树的可视化

# 把这棵树clf进行图像化，让人看着更加清晰
dot_data = tree.export_graphviz(regressor,  # 决策树本身
                                feature_names=boston.feature_names,  # 特征名称
                                filled=True,  # 给图形填充颜色
                                rounded=True  # 图形的节点是圆角矩形
                                )
graph = graphviz.Source(dot_data)
graph.view(filename='reg.png', cleanup=True)

如下图：

将上面的代码连起来就是完整实例代码咯。