如何在sklearn决策树中显示特征名称?

时间:2021-03-09 18:07:14

标签: python pandas machine-learning scikit-learn decision-tree

我目前有一个决策树,将特征名称显示为 X[index],即 X[0], X[1], X[2]

from sklearn import tree
from sklearn.tree import DecisionTreeClassifier

dt = DecisionTreeClassifier()
dt.fit(X_train, y_train)

 
# plot tree
plt.figure(figsize=(20,16))# set plot size (denoted in inches)
tree.plot_tree(dt,fontsize=10)

我希望将这些 X[featureNumber] 替换为实际的功能名称。

因此,与其显示 X[0],我希望它显示由 X.columns.values[0] 返回的功能名称(我不知道此代码是否正确)。

我也知道使用 graphviz 有一种简单的方法可以做到这一点,但由于某种原因我无法在 Jupiter 中运行 graphviz,所以我正在寻找一种不用的方法。

当前决策树的照片:

1 个答案:

答案 0 :(得分:2)

这在 the documentation 中有解释:

<块引用>

sklearn.tree.plot_tree(decision_tree, *, max_depth=None, feature_names=None, class_names=None, label='all', filled=False, impurity=True, node_ids=False, proportion=False, rotate='deprecated', rounded=False, precision=3, ax=None, fontsize=None)

feature_names:字符串列表,默认=None

每个功能的名称。如果 None,将使用通用名称(“X[0]”、“X[1]”、...)。