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

时间:2017-06-29 09:31:15

标签: python scikit-learn decision-tree

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

enter image description here

为了让人们更容易解释,我只想保留决定和课程。 我在哪里可以修改代码来执行此操作?

谢谢。

1 个答案:

答案 0 :(得分:2)

根据documentation,无法在框内设置附加信息。您可能隐含省略的唯一事情是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)

以下是修改前后的图像:

enter image description here enter image description here

在你的情况下,框中似乎有更多参数,所以你可能想稍微调整一下代码。

我希望有所帮助!