如何在scikit中为决策树的输出添加功能名称了解

时间:2015-09-11 18:50:28

标签: python scikit-learn

我正在使用scikit-learn构建决策树,然后想要生成树的pdf。我输出树的工作流程大致如下。

vec = DictVectorizer()
data_vectorized = vec.fit_transform(data)
vec.get_feature_names() #Shows feature names

clf = tree.DecisionTreeClassifier()
clf = clf.fit(data_vectorized, Labels)

from sklearn.externals.six import StringIO  
import pydot 
dot_data = StringIO() 
tree.export_graphviz(clf, out_file=dot_data) 
graph = pydot.graph_from_dot_data(dot_data.getvalue()) 
graph.write_pdf("tree.pdf") 

然而,这个pdf的每个节点都显示了对某些i使用data_vectorized [i]的比较。当data_vectorized很大且稀疏时,这很难理解。

如何让它显示该功能的名称?

此图显示了您获得的示例(当要素位于变量X中时。例如,我希望将X [2]替换为要素的名称。

enter image description here

1 个答案:

答案 0 :(得分:3)

尝试将导出更改为:

tree.export_graphviz(clf, out_file=dot_data, feature_names=vec.get_feature_names())