决策树可视化

决策树可视化方法

例1.

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")

例2.

from sklearn.tree import export_graphviz
import pydotplus
from IPython.display import Image

rf_small = RandomForestRegressor(n_estimators=10, max_depth = 3, random_state=42)
rf_small.fit(train_features, train_labels)

# 提取一颗树
tree_small = rf_small.estimators_[5]

# 保存
export_graphviz(tree_small, out_file = 'small_tree.dot', feature_names = features_columns, rounded = True, precision = 1)

graph = pydotplus.graph_from_dot_file('small_tree.dot')

display(Image(graph.create_png()))

graph.write_png('small_tree.png');