Python中决策树的模型可视化方法

2023-04-14 00:00:00 模型 方法 可视化

Python中可以使用Graphviz库来可视化决策树模型。以下是具体的步骤：

1.安装Graphviz库：可以使用pip安装，命令如下：

pip install graphviz

2.导入需要的库

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier, export_graphviz
from sklearn.model_selection import train_test_split
from graphviz import Source

3.加载数据集

这里以鸢尾花数据集为例进行演示

iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, random_state=0)

4.训练模型

clf = DecisionTreeClassifier(max_depth=2, random_state=0)
clf.fit(X_train, y_train)

5.生成图形化模型

使用export_graphviz函数生成.dot格式的决策树模型表述，然后将其转换为图形化模型。其中，需要指定一些参数，如特征名称、目标名称等。

dot_data = export_graphviz(clf, out_file=None,
                           feature_names=iris.feature_names,
                           class_names=iris.target_names,
                           filled=True, rounded=True,
                           special_characters=True)
graph = Source(dot_data)
graph.view()

此时，系统会自动弹出一个窗口，显示决策树模型。

可以看到，我们生成了一个带有中文字符的决策树模型。其中，左侧的箭头表示指向小于等于给定值的情况，右侧的箭头表示指向大于给定值的情况。叶子节点则表示决策结果。matplotlib也可用于决策树的绘制，但其不如graphviz方便。

相关文章