决策树可视化

决策树可视化方法

from sklearn.model_selection import train_test_split
from sklearn.datasets.california_housing import fetch_california_housing
import pydotplus
from IPython.display import Image

housing = fetch_california_housing()

# 数据集划分
data_train, data_test, target_train, target_test = train_test_split(housing.data, 
                                                                    housing.target,
                                                                   test_size=0.1,
                                                                   random_state=42)
# 决策树
dtr = tree.DecisionTreeRegressor(random_state=42)
dtr.fit(data_train, target_train)

dot_dtr = tree.export_graphviz(dtr, 
                               out_file=None,
                               feature_names=housing.feature_names,
                               filled=True,
                               impurity=False,
                               rounded=True)


graph = pydotplus.graph_from_dot_data(dot_data)

graph.get_nodes()[7].set_fillcolor("#ADD2AA")

# 可视化决策树
Image(graph.create_png())

# 将图像保存
graph.write_png("dtr.png")