Python - Graphviz - 删除 DecisionTreeClassifier 节点上的图例

Dio*_*ong 3 python decision-tree scikit-learn

我有一个来自 sklearn 的决策树分类器,我使用 pydotplus 来展示它。但是,当我的演示文稿的每个节点上都有很多信息(熵、样本和值)时,我真的不喜欢。

在此处输入图片说明

为了向人们解释它更容易,我只想保留决定和课程。我在哪里可以修改代码来做到这一点?

谢谢你。

E.Z*_*.Z. 7

根据文档,不可能避免在框内设置附加信息。唯一可以隐式省略的是impurity参数。

但是,我以另一种显式的方式完成了它,这有点歪。首先,我保存.dot文件,将杂质设置为False。然后,我打开它并将其转换为字符串格式。我使用正则表达式减去冗余标签并重新保存。

代码是这样的:

import pydotplus  # pydot library: install it via pip install pydot
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz
from sklearn.datasets import load_iris

data = load_iris()
clf = DecisionTreeClassifier()
clf.fit(data.data, data.target)

export_graphviz(clf, out_file='tree.dot', impurity=False, class_names=True)

PATH = '/path/to/dotfile/tree.dot'
f = pydot.graph_from_dot_file(PATH).to_string()
f = re.sub('(\\\\nsamples = [0-9]+)(\\\\nvalue = \[[0-9]+, [0-9]+, [0-9]+\])', '', f)
f = re.sub('(samples = [0-9]+)(\\\\nvalue = \[[0-9]+, [0-9]+, [0-9]+\])\\\\n', '', f)

with open('tree_modified.dot', 'w') as file:
    file.write(f)
Run Code Online (Sandbox Code Playgroud)

以下是修改前后的图片:

在此处输入图片说明在此处输入图片说明

在您的情况下,框中似乎有更多参数,因此您可能需要稍微调整代码。

我希望这有帮助!